fix(group): emit consecutive MI integers 0..N-1

src/lib/commands/dedup.rs

@@ -177,6 +177,10 @@ pub struct ProcessedDedupGroup {
     pub dedup_metrics: DedupMetrics,
     /// Total input records processed (for progress tracking).
     pub input_record_count: u64,
+    /// Number of distinct numeric molecule IDs assigned in this group. See
+    /// [`fgumi_lib::grouper::ProcessedPositionGroup::distinct_mi_count`] for
+    /// rationale; ensures emitted MI tag integers are consecutive 0..N-1.
+    pub distinct_mi_count: u64,
 }
 
 impl BatchWeight for ProcessedDedupGroup {
@@ -883,6 +887,7 @@ fn process_position_group(
             family_sizes: AHashMap::new(),
             dedup_metrics,
             input_record_count,
+            distinct_mi_count: 0,
         });
     }
 
@@ -999,7 +1004,18 @@ fn process_position_group(
 
     dedup_metrics.unique_reads = dedup_metrics.total_reads - dedup_metrics.duplicate_reads;
 
-    Ok(ProcessedDedupGroup { templates, family_sizes, dedup_metrics, input_record_count })
+    // Compute distinct numeric MoleculeId count for the global MI counter block
+    // size; see issue #269 / ProcessedPositionGroup::distinct_mi_count rationale.
+    let distinct_mi_count: u64 =
+        templates.iter().filter_map(|t| t.mi.id()).max().map(|max_id| max_id + 1).unwrap_or(0);
+
+    Ok(ProcessedDedupGroup {
+        templates,
+        family_sizes,
+        dedup_metrics,
+        input_record_count,
+        distinct_mi_count,
+    })
 }
 
 //////////////////////////////////////////////////////////////////////////////
@@ -1277,11 +1293,12 @@ impl Command for MarkDuplicates {
 
                 let input_record_count = processed.input_record_count;
 
-                // Assign contiguous base_mi from serial counter
-                let base_mi = next_mi_base.fetch_add(
-                    processed.templates.len() as u64,
-                    std::sync::atomic::Ordering::Relaxed,
-                );
+                // Assign contiguous base_mi from the global counter. Advance by the
+                // number of distinct numeric MoleculeId IDs actually assigned in this
+                // group, not the template count, so the emitted MI integers are
+                // consecutive 0..N-1 (see issue #269).
+                let base_mi = next_mi_base
+                    .fetch_add(processed.distinct_mi_count, std::sync::atomic::Ordering::Relaxed);
                 let assign_tag = Tag::from(assign_tag_bytes);
 
                 // Pre-allocate output buffer: ~2 records/template × ~400 bytes/record
@@ -2885,6 +2902,7 @@ mod tests {
             family_sizes: AHashMap::new(),
             dedup_metrics: DedupMetrics::default(),
             input_record_count: 1,
+            distinct_mi_count: 0,
         };
 
         let estimate = batch.estimate_heap_size();

src/lib/commands/group.rs

@@ -1117,9 +1117,11 @@ impl Command for GroupReadsByUmi {
         #[cfg(feature = "memory-debug")]
         let short_circuit_compress = short_circuit == "compress";
 
-        // Counter for contiguous MI assignment (incremented in the serial serialize step).
-        // AtomicU64 satisfies the Fn + Sync bound; Relaxed ordering is fine because
-        // the serialize step is serial and ordered.
+        // Counter for contiguous MI assignment. The serialize step is parallel, so each
+        // thread reserves its own block via fetch_add of the per-group
+        // `distinct_mi_count`. AtomicU64 satisfies the Fn + Sync bound; Relaxed ordering
+        // is fine because the counter is never read out-of-band (only the values produced
+        // by fetch_add are used, and each serialize thread uses only its own reservation).
         let next_mi_base = std::sync::atomic::AtomicU64::new(0);
 
         // Run the 7-step unified pipeline with the already-opened reader (supports streaming)
@@ -1145,6 +1147,7 @@ impl Command for GroupReadsByUmi {
                         family_sizes: AHashMap::new(),
                         filter_metrics: FilterMetrics::new(),
                         input_record_count,
+                        distinct_mi_count: 0,
                     });
                 }
                 let mut filter_metrics = FilterMetrics::new();
@@ -1229,6 +1232,7 @@ impl Command for GroupReadsByUmi {
                         family_sizes: AHashMap::new(),
                         filter_metrics,
                         input_record_count,
+                        distinct_mi_count: 0,
                     });
                 }
 
@@ -1276,9 +1280,22 @@ impl Command for GroupReadsByUmi {
                         family_sizes: AHashMap::new(),
                         filter_metrics,
                         input_record_count,
+                        distinct_mi_count: 0,
                     });
                 }
 
+                // Compute the number of distinct numeric molecule IDs assigned in this group.
+                // Assigners hand out numeric IDs 0, 1, 2, ... contiguously, so `max(id) + 1`
+                // equals the count of distinct IDs used. This can be less than
+                // `templates.len()` when multiple templates share a UMI family (and hence a
+                // MoleculeId) or when paired A/B variants share the same numeric id.
+                let distinct_mi_count: u64 = templates
+                    .iter()
+                    .filter_map(|t| t.mi.id())
+                    .max()
+                    .map(|max_id| max_id + 1)
+                    .unwrap_or(0);
+
                 // Sort templates directly by (MI index, name) - avoids Vec<Vec<Template>> allocation
                 templates.sort_by(|a, b| {
                     let a_idx = a.mi.to_vec_index();
@@ -1327,6 +1344,7 @@ impl Command for GroupReadsByUmi {
                     family_sizes,
                     filter_metrics,
                     input_record_count,
+                    distinct_mi_count,
                 })
             },
             // serialize_fn: Serialize records + collect metrics (serial, ordered)
@@ -1360,9 +1378,14 @@ impl Command for GroupReadsByUmi {
                 // Save input record count for progress tracking
                 let input_record_count = processed.input_record_count;
 
-                // Assign contiguous base_mi from serial counter (template count, not record count)
+                // Assign contiguous base_mi from the global counter. Advance by the
+                // number of distinct numeric MoleculeId IDs actually assigned in this
+                // group (not the template count), because multiple templates can share
+                // the same MoleculeId and because PairedA/PairedB share a numeric id.
+                // This ensures emitted MI integers are consecutive 0..N-1, matching
+                // fgbio's `GroupReadsByUmi` (see issue #269).
                 let base_mi = next_mi_base.fetch_add(
-                    processed.templates.len() as u64,
+                    processed.distinct_mi_count,
                     std::sync::atomic::Ordering::Relaxed,
                 );
 
@@ -1736,9 +1759,17 @@ impl GroupReadsByUmi {
             }
         }
 
-        // Write templates (already sorted by MI, then by name)
+        // Write templates (already sorted by MI, then by name).
+        //
+        // Advance the global MI counter by the number of distinct numeric MoleculeId
+        // IDs actually assigned in this group, not by the template count, so the
+        // emitted MI integers are consecutive 0..N-1 across all position groups
+        // (matching fgbio's `GroupReadsByUmi`; see issue #269). Assigners hand out
+        // numeric IDs 0, 1, 2, ... contiguously, so `max(id) + 1` equals the count.
+        let distinct_mi_count: u64 =
+            templates.iter().filter_map(|t| t.mi.id()).max().map(|max_id| max_id + 1).unwrap_or(0);
         let base_mi = *next_mi_base;
-        *next_mi_base += templates.len() as u64;
+        *next_mi_base += distinct_mi_count;
         let assign_tag = Tag::from(assign_tag_bytes);
 
         for template in templates {

src/lib/grouper.rs

@@ -368,6 +368,16 @@ pub struct ProcessedPositionGroup {
     pub filter_metrics: FilterMetrics,
     /// Total input records processed (for progress tracking).
     pub input_record_count: u64,
+    /// Number of distinct numeric molecule IDs assigned in this group (i.e., the
+    /// size of the block the serialize step must reserve from the global MI
+    /// counter). Equals `max(MoleculeId::id()) + 1` across assigned templates,
+    /// or `0` when no templates have an assigned MI. This can be strictly less
+    /// than `templates.len()` because templates in the same UMI family share
+    /// a `MoleculeId`, and because `PairedA(id)` / `PairedB(id)` share the same
+    /// numeric `id`. Using this value for the block size (rather than the
+    /// template count) ensures the emitted MI tag integers are consecutive
+    /// `0..N-1`, matching fgbio's `GroupReadsByUmi`.
+    pub distinct_mi_count: u64,
 }
 
 // ============================================================================

tests/integration/test_group_command.rs

@@ -2,8 +2,13 @@
 
 use fgoxide::io::DelimFile;
 use fgumi_lib::metrics::UmiGroupingMetrics;
+use fgumi_lib::sam::builder::RecordBuilder;
 use noodles::bam;
 use noodles::sam::alignment::io::Write as AlignmentWrite;
+use noodles::sam::alignment::record::data::field::Tag;
+use noodles::sam::alignment::record_buf::RecordBuf;
+use noodles::sam::alignment::record_buf::data::field::value::Value as DataValue;
+use std::collections::HashSet;
 use std::fs;
 use std::path::PathBuf;
 use std::process::Command;
@@ -147,3 +152,185 @@ fn create_test_bam_with_n_umis(path: &PathBuf) {
 
     writer.try_finish().expect("Failed to finish BAM");
 }
+
+/// Create a group of single-end reads at a specific position with the same UMI.
+///
+/// All reads share the same reference, position, and UMI, so they form a single
+/// UMI family within their position group.
+fn create_umi_family_at(base_name: &str, umi: &str, count: usize, start: usize) -> Vec<RecordBuf> {
+    (0..count)
+        .map(|i| {
+            RecordBuilder::new()
+                .name(&format!("{base_name}_{i}"))
+                .sequence("ACGTACGT")
+                .qualities(&[30; 8])
+                .reference_sequence_id(0)
+                .alignment_start(start)
+                .mapping_quality(60)
+                .cigar("8M")
+                .tag("RX", umi)
+                .build()
+        })
+        .collect()
+}
+
+/// Regression test for issue #269: the `group` command should emit consecutive
+/// MI integer values in the range `[0, N-1]`, where `N` is the number of distinct
+/// UMI families. fgbio's `GroupReadsByUmi` has this property; fgumi should too.
+///
+/// The input is constructed so that several position groups each contain more
+/// templates than distinct UMI families, forcing the old per-thread block
+/// allocator (which advanced by `templates.len()`) to leave gaps in the
+/// emitted MI integer space.
+#[test]
+fn test_group_command_assigns_consecutive_mi_values() {
+    let temp_dir = TempDir::new().expect("Failed to create temp dir");
+    let input_bam = temp_dir.path().join("input.bam");
+    let output_bam = temp_dir.path().join("output.bam");
+
+    // Build the input BAM: three position groups, each with templates.len() > distinct UMIs.
+    //
+    // Group 1 (position 100):  4 reads, 1 UMI       -> 1 distinct MI, 3 wasted slots
+    // Group 2 (position 500):  5 reads, 2 UMIs      -> 2 distinct MIs, 3 wasted slots
+    // Group 3 (position 1000): 3 reads, 1 UMI       -> 1 distinct MI, 2 wasted slots
+    // Total distinct MIs expected: 4
+    let mut records: Vec<RecordBuf> = Vec::new();
+    records.extend(create_umi_family_at("g1", "AAAAAAAA", 4, 100));
+    records.extend(create_umi_family_at("g2a", "CCCCCCCC", 3, 500));
+    records.extend(create_umi_family_at("g2b", "GGGGGGGG", 2, 500));
+    records.extend(create_umi_family_at("g3", "TTTTTTTT", 3, 1000));
+
+    let header = create_minimal_header("chr1", 10000);
+    let mut writer =
+        bam::io::Writer::new(fs::File::create(&input_bam).expect("Failed to create BAM file"));
+    writer.write_header(&header).expect("Failed to write header");
+    for record in &records {
+        writer.write_alignment_record(&header, record).expect("Failed to write record");
+    }
+    writer.try_finish().expect("Failed to finish BAM");
+
+    let status = Command::new(env!("CARGO_BIN_EXE_fgumi"))
+        .args([
+            "group",
+            "--input",
+            input_bam.to_str().unwrap(),
+            "--output",
+            output_bam.to_str().unwrap(),
+            "--strategy",
+            "identity",
+            "--edits",
+            "0",
+            "--compression-level",
+            "1",
+        ])
+        .status()
+        .expect("Failed to run group command");
+    assert!(status.success(), "Group command failed");
+    assert!(output_bam.exists(), "Output BAM not created");
+
+    // Collect the set of MI tag values written to the output.
+    let mut reader = bam::io::Reader::new(fs::File::open(&output_bam).unwrap());
+    let out_header = reader.read_header().expect("read output header");
+    let mi_tag = Tag::from(*b"MI");
+
+    let mut mi_values: HashSet<u64> = HashSet::new();
+    for result in reader.record_bufs(&out_header) {
+        let record = result.expect("read record");
+        if let Some(DataValue::String(mi)) = record.data().get(&mi_tag) {
+            let s = std::str::from_utf8(mi).expect("MI tag is utf8");
+            let value: u64 = s.parse().expect("MI tag is a non-negative integer");
+            mi_values.insert(value);
+        }
+    }
+
+    assert_eq!(mi_values.len(), 4, "expected 4 distinct MI values (one per UMI family)");
+
+    let max_mi = *mi_values.iter().max().expect("at least one MI value");
+    let expected_max = (mi_values.len() as u64) - 1;
+    assert_eq!(
+        max_mi, expected_max,
+        "MI values should be consecutive 0..N-1; got {mi_values:?} (max {max_mi}, expected {expected_max})"
+    );
+}
+
+/// Regression test for issue #269 covering the multi-threaded pipeline path.
+///
+/// Same invariant (consecutive MI integers 0..N-1) but forces the unified BAM
+/// pipeline by running with multiple threads. The counter in that path is an
+/// `AtomicU64` shared across parallel serialize workers, so this test guards
+/// against the per-worker block-allocation form of the bug as well.
+#[test]
+fn test_group_command_assigns_consecutive_mi_values_multi_threaded() {
+    let temp_dir = TempDir::new().expect("Failed to create temp dir");
+    let input_bam = temp_dir.path().join("input.bam");
+    let output_bam = temp_dir.path().join("output.bam");
+
+    // Build a larger input with many position groups so multiple worker threads
+    // can reserve MI blocks in parallel. Each position group has more templates
+    // than distinct UMI families so the old per-template block size would leak
+    // gaps into the emitted MI space.
+    let mut records: Vec<RecordBuf> = Vec::new();
+    let num_groups = 20;
+    for g in 0..num_groups {
+        let start = 100 + (g * 500);
+        // Two UMI families per position group, 4 and 3 reads respectively.
+        records.extend(create_umi_family_at(&format!("g{g}_a"), "AAAAAAAA", 4, start));
+        records.extend(create_umi_family_at(&format!("g{g}_b"), "CCCCCCCC", 3, start));
+    }
+
+    let header = create_minimal_header("chr1", 1_000_000);
+    let mut writer =
+        bam::io::Writer::new(fs::File::create(&input_bam).expect("Failed to create BAM file"));
+    writer.write_header(&header).expect("Failed to write header");
+    for record in &records {
+        writer.write_alignment_record(&header, record).expect("Failed to write record");
+    }
+    writer.try_finish().expect("Failed to finish BAM");
+
+    let status = Command::new(env!("CARGO_BIN_EXE_fgumi"))
+        .args([
+            "group",
+            "--input",
+            input_bam.to_str().unwrap(),
+            "--output",
+            output_bam.to_str().unwrap(),
+            "--strategy",
+            "identity",
+            "--edits",
+            "0",
+            "--threads",
+            "4",
+            "--compression-level",
+            "1",
+        ])
+        .status()
+        .expect("Failed to run group command");
+    assert!(status.success(), "Group command failed");
+
+    let mut reader = bam::io::Reader::new(fs::File::open(&output_bam).unwrap());
+    let out_header = reader.read_header().expect("read output header");
+    let mi_tag = Tag::from(*b"MI");
+
+    let mut mi_values: HashSet<u64> = HashSet::new();
+    for result in reader.record_bufs(&out_header) {
+        let record = result.expect("read record");
+        if let Some(DataValue::String(mi)) = record.data().get(&mi_tag) {
+            let s = std::str::from_utf8(mi).expect("MI tag is utf8");
+            let value: u64 = s.parse().expect("MI tag is a non-negative integer");
+            mi_values.insert(value);
+        }
+    }
+
+    let expected_distinct = num_groups * 2;
+    assert_eq!(
+        mi_values.len(),
+        expected_distinct,
+        "expected {expected_distinct} distinct MI values (two per position group)"
+    );
+    let max_mi = *mi_values.iter().max().expect("at least one MI value");
+    let expected_max = (mi_values.len() as u64) - 1;
+    assert_eq!(
+        max_mi, expected_max,
+        "MI values should be consecutive 0..N-1 under multi-threaded pipeline; got max {max_mi}, expected {expected_max}"
+    );
+}