perf(sort): optimize coordinate sort with radix sort, memory scaling, and batched I/O

[nh13]

Summary

Performance optimizations for coordinate sorting that bring fgumi in line with or exceeding samtools:

• Radix sort: Add LSD radix sort for RecordRef with O(n×k) complexity vs O(n log n) comparison sort
• Memory scaling: 768MB per thread by default (was fixed 512MB total), matching samtools behavior
• Batched reading: 256 records per channel send in ReadAheadReader (was 1), reducing sync overhead ~256x

Bug Fixes

• Fix non-stable k-way merge by adding chunk_idx as tie-breaker in heap comparators
• Add temp file consolidation (max 64 files, merges oldest when exceeded)

Infrastructure

• Add raw BGZF block reader for future direct reading optimization
• Add comprehensive radix sort unit tests

Benchmark Results

Tested on 1.9GB BAM file (59M records), 8 threads:

Tool	Coordinate Sort	vs Baseline
fgumi (optimized)	24.7s	-18%
fgumi (baseline)	30.0s	-
samtools	29.5s	-

fgumi now outperforms samtools by 16% on coordinate sorting.

Test plan

• All 2,027 existing tests pass
• New radix sort unit tests pass
• Benchmarked against samtools on large files
• Output verified correct via fgumi compare bams --mode content

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds per-thread memory scaling for sorting (new memory_per_thread: bool, default true) and increases default max_memory to 768M. Replaces single-threaded sorts with radix and parallel radix sorts for RecordRef and TemplateRecordRef and exposes new radix sort utilities. Adds raw_bam_reader module with RawBamRecordReader and BatchedRawBamReader for streamed BGZF/raw BAM access. Enhances RawExternalSorter with temp-file consolidation (max_temp_files, default 64) and stable merge tie-breaking. Converts read-ahead to batched I/O (batched record delivery, new batching APIs and constructors).

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main changes: radix sort optimization, memory scaling per thread, and batched I/O—all core performance improvements in this PR.
Description check	✅ Passed	The description is directly related to the changeset, detailing radix sort implementation, memory scaling, batched reading, bug fixes, and providing benchmark results that validate the changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch perf/sort-optimizations

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/commands/sort.rs (1)

210-243: ⚠️ Potential issue | 🟡 Minor

Clamp/validate threads consistently with memory_per_thread.
effective_memory uses threads.max(1) but the sorter still gets the raw self.threads. If 0 slips through, logs/memory math diverge and downstream may panic.

Proposed fix

-        let effective_memory = if self.memory_per_thread {
-            self.max_memory.saturating_mul(self.threads.max(1))
-        } else {
-            self.max_memory
-        };
+        let threads = self.threads.max(1);
+        let effective_memory = if self.memory_per_thread {
+            self.max_memory.saturating_mul(threads)
+        } else {
+            self.max_memory
+        };
@@
-        info!("Threads: {}", self.threads);
+        info!("Threads: {}", threads);
@@
-        let mut sorter = RawExternalSorter::new(self.order.into())
-            .memory_limit(effective_memory)
-            .threads(self.threads)
+        let mut sorter = RawExternalSorter::new(self.order.into())
+            .memory_limit(effective_memory)
+            .threads(threads)

🤖 Fix all issues with AI agents

In `@src/lib/sort/raw_bam_reader.rs`:
- Around line 172-209: The docstring for next_record() is incorrect: it claims
the returned Vec<u8> includes the 4-byte block_size prefix but the
implementation in next_record() (which reads block_size and then sets record =
self.decompressed[self.position + 4..self.position + total_size].to_vec())
intentionally strips those 4 bytes and returns only the record payload; update
the doc/comment to state clearly that the returned bytes exclude the 4-byte
block_size prefix (i.e., they are the BAM record payload only) and mention that
total_size = 4 + block_size is used for validation, or alternatively change the
code to include the 4-byte prefix if you prefer keeping the original doc—make
the doc and the behavior of next_record() consistent.
- Around line 124-166: skip_header() calls to ensure_bytes() currently use the ?
operator and ignore the bool return (false on truncated input), leading to
panics when slicing; update skip_header() to check the bool result of
ensure_bytes(...) after each call (for header text, ref name, etc.) and if it is
false return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "truncated BAM
header")); similarly, update read_u32() to treat a false from ensure_bytes(...)
as an UnexpectedEof error instead of proceeding; follow the same pattern used in
next_record() (checking the boolean and returning UnexpectedEof) so all
slices/indexing in skip_header() and read_u32() are guarded against truncated
input.

In `@src/lib/sort/raw.rs`:
- Around line 511-618: The consolidation logic in maybe_consolidate_temp_files
can leave merge_count == 0 and breaks stability by appending the merged file as
newest; fix by computing merge_count = std::cmp::max(1, chunk_files.len()/2) so
we always reduce files when needed, preserve original chunk order by capturing
each file's original index when building readers (use
files_to_merge.iter().enumerate() and store original_chunk_idx in HeapEntry),
use that original_chunk_idx as the tie-breaker in the lt closure instead of
reader_idx, and insert the merged_path back into chunk_files at the position of
the removed oldest files (e.g., chunk_files.insert(0, merged_path) or insert at
the drained range start) rather than push so equal-key records keep the original
relative order.

In `@src/lib/sort/read_ahead.rs`:
- Around line 61-83: with_buffer_size currently always uses BATCH_SIZE for
batch_size which overshoots when buffer_size < BATCH_SIZE; change it to derive
batch_size from buffer_size and guard against zero: compute batch_size =
buffer_size.min(BATCH_SIZE).max(1) and channel_buffer = ((buffer_size +
batch_size - 1) / batch_size).max(1) (ceil division with a 0 guard), then call
Self::with_batch_size(reader, batch_size, channel_buffer); reference symbols:
with_buffer_size, with_batch_size, BATCH_SIZE, channel_buffer, batch_size.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@src/lib/sort/read_ahead.rs`:
- Around line 61-75: The receiver must be closed before joining the reader
thread to avoid a drop-time deadlock; change the ReadAheadReader.receiver field
to Option<Receiver<...>> and update constructors such as with_batch_size to
store Some(rx), then modify Drop for ReadAheadReader to take the receiver (e.g.,
self.receiver.take()) and drop it before calling join on the thread handle (use
handle.take() and join()); ensure reader_thread and any code that accesses
receiver account for Option by unwrapping or using take() where appropriate so
the channel is closed prior to join.

🧹 Nitpick comments (2)

src/lib/sort/raw_bam_reader.rs (2)

313-317: Consider validating batch_size > 0.

A batch_size of 0 causes next_batch() to always return None, appearing as immediate EOF. Consider:

assert!(batch_size > 0, "batch_size must be positive");

Or use NonZero<usize> in the signature.

---

342-349: Test coverage is minimal.

Only a constant check exists. When this module is integrated into the sort pipeline, add tests for:

• Valid BAM parsing
• Truncated input handling
• Batch boundary behavior

Acceptable for infrastructure code not yet in use.