Repository navigation
perf(group): cache UMI position to drop redundant aux scan (#334) - #343
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (8)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (5)
📝 WalkthroughWalkthroughAdds decode-time discovery of Z-type UMI aux-tag offsets in raw-bam, stores offset+length on DecodedRecord when configured, propagates that cache into Template::from_decoded_records and exposes Template::cached_umi(), and updates grouping to prefer cached slices (with tests and small rewrap of re-exports). 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #343 +/- ##
==========================================
+ Coverage 90.87% 90.91% +0.03%
==========================================
Files 78 78
Lines 50775 51017 +242
==========================================
+ Hits 46143 46381 +238
- Misses 4632 4636 +4 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/lib/unified_pipeline/bam.rs (1)
4853-4874: ⚡ Quick winAdd a non-
ZRXcase here.
cache_umi_position()documents the same fallback for missing and non-string tags, but this only locks the missing-tag branch. One non-ZRXrecord would pin the other half of the contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/lib/unified_pipeline/bam.rs` around lines 4853 - 4874, Add a new unit test alongside test_cache_umi_position_leaves_cache_unset_when_tag_missing that constructs a record with an RX tag present but not of type 'Z' (e.g., an integer or binary tag value), then call cache_umi_position(&mut decoded, *SamTag::RX) and assert that decoded.cached_umi().is_none() and decoded.cached_umi_position().0 == DecodedRecord::UMI_OFFSET_UNCACHED; this verifies cache_umi_position correctly treats non-'Z' RX tags the same as missing tags (refer to cache_umi_position, DecodedRecord, and SamTag::RX to locate the logic).src/lib/commands/group.rs (1)
946-950: ⚡ Quick winSkip decode-time UMI caching in
--no-umimode.
assign_umi_groups_implnever reads a UMI whenno_umiis set, so this still pays the aux scan in Decode with no consumer. Gatewith_umi_tag(raw_tag)behind!self.no_umito keep that mode on the cheaper path.Proposed change
let library_index = LibraryIndex::from_header(&header); - // Cache the UMI tag's value position on each DecodedRecord so the Process - // step's UMI-assignment pass can slice the value without re-scanning aux - // data (issue `#334`). - pipeline_config.group_key_config = - Some(GroupKeyConfig::new(library_index, cell_tag).with_umi_tag(raw_tag)); + // Cache the UMI tag's value position on each DecodedRecord so the Process + // step's UMI-assignment pass can slice the value without re-scanning aux + // data (issue `#334`). + let mut group_key_config = GroupKeyConfig::new(library_index, cell_tag); + if !self.no_umi { + group_key_config = group_key_config.with_umi_tag(raw_tag); + } + pipeline_config.group_key_config = Some(group_key_config);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/lib/commands/group.rs` around lines 946 - 950, The code always adds the UMI tag into pipeline_config.group_key_config via GroupKeyConfig::new(...).with_umi_tag(raw_tag) even in --no-umi mode; change it to only call with_umi_tag(raw_tag) when self.no_umi is false. Concretely, create the GroupKeyConfig with GroupKeyConfig::new(library_index, cell_tag), then if !self.no_umi call .with_umi_tag(raw_tag) on that config before assigning pipeline_config.group_key_config = Some(config); otherwise assign the config without the UMI tag.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/lib/commands/group.rs`:
- Around line 946-950: The code always adds the UMI tag into
pipeline_config.group_key_config via
GroupKeyConfig::new(...).with_umi_tag(raw_tag) even in --no-umi mode; change it
to only call with_umi_tag(raw_tag) when self.no_umi is false. Concretely, create
the GroupKeyConfig with GroupKeyConfig::new(library_index, cell_tag), then if
!self.no_umi call .with_umi_tag(raw_tag) on that config before assigning
pipeline_config.group_key_config = Some(config); otherwise assign the config
without the UMI tag.
In `@src/lib/unified_pipeline/bam.rs`:
- Around line 4853-4874: Add a new unit test alongside
test_cache_umi_position_leaves_cache_unset_when_tag_missing that constructs a
record with an RX tag present but not of type 'Z' (e.g., an integer or binary
tag value), then call cache_umi_position(&mut decoded, *SamTag::RX) and assert
that decoded.cached_umi().is_none() and decoded.cached_umi_position().0 ==
DecodedRecord::UMI_OFFSET_UNCACHED; this verifies cache_umi_position correctly
treats non-'Z' RX tags the same as missing tags (refer to cache_umi_position,
DecodedRecord, and SamTag::RX to locate the logic).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0c5e9666-5b93-4cc1-93fa-2d5e63a580c0
📒 Files selected for processing (8)
crates/fgumi-raw-bam/src/lib.rscrates/fgumi-raw-bam/src/tags.rssrc/lib/commands/dedup.rssrc/lib/commands/group.rssrc/lib/grouper.rssrc/lib/template.rssrc/lib/unified_pipeline/bam.rssrc/lib/unified_pipeline/base.rs
…te aux scan The parallel Decode step already scans each record's aux data for RG, CB, and MC. After this change it also records the UMI tag's value position (`umi_value_offset: u32`, `umi_value_len: u16`) on the `DecodedRecord` when `GroupKeyConfig::umi_tag` is set. Template construction propagates the cached position to a new `Template::cached_umi_position`, and the UMI-assignment step in `fgumi group` slices the value through `Template::cached_umi()` instead of calling `find_string_tag` a second time. A new `fgumi_raw_bam::find_string_tag_position` helper returns `Option<(u32, u16)>` so the cached offset can be expressed without allocating or copying. Templates built outside the Decode path (tests, non-group callers) leave the cache `None` and continue to use the existing fallback. Resolves #334.
1a733b0 to
c32fa5d
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
Resolves #334. Two improvements landed together because the smaller one revealed the larger one: 1. fgumi dedup now opts in to the per-record UMI position cache (#343 infrastructure) and reads via template.cached_umi() in its UMI- assignment loop with a fallback to find_string_tag. Skipped in --no-umi mode where the assignment site emits String::new(). 2. The cache populate, previously a separate second walk of aux data via cache_umi_position, is now folded into extract_aux_string_tags so a single linear walk captures RG, cell, MC, and the UMI position together. cache_umi_position is deleted; aux_data_offset_from_record is now called exactly once per record on the decode hot path. Both improvements affect every command that consumes the UMI cache via the parallel pipeline (currently fgumi group from #343 and fgumi dedup from this commit). Mechanics: - AuxStringTags gains umi_position: Option<(u32, u16)> (aux-relative value offset + length, NUL excluded). - extract_aux_string_tags takes umi_tag: Option<[u8; 2]>; the inner loop captures UMI position in the same branch ladder as RG/cell/MC; early-exit mask widens from 7 to 15 when umi_tag is requested. - compute_group_key_from_raw returns (GroupKey, Option<(u32, u16)>); decode_records converts aux-relative to record-relative and stores via decoded.set_cached_umi(...). - Single-threaded group fast path (commands/group.rs:1408) passes None for umi_tag — it doesn't populate the cache and didn't before. Measurements (test workload, --threads 4, warm-cache tricord, profiling-profile binary): fgumi dedup (3-run median, runs within 0.02s): parent (PR343) this commit delta wall (s) 17.82 17.67 -0.8% cpu_time (s) 68.83 67.84 -1.4% max_rss (MB) 1435 1435 flat fgumi group: no detectable wall improvement (8-run min matches parent at 17.66s; per-record CPU saving below the noise floor at this workload size). Improvement remains architecturally real. Output BAM byte-identical pre/post on dedup across 53.1M records.
Resolves #334. Two improvements landed together because the smaller one revealed the larger one: 1. fgumi dedup now opts in to the per-record UMI position cache (#343 infrastructure) and reads via template.cached_umi() in its UMI- assignment loop with a fallback to find_string_tag. Skipped in --no-umi mode where the assignment site emits String::new(). 2. The cache populate, previously a separate second walk of aux data via cache_umi_position, is now folded into extract_aux_string_tags so a single linear walk captures RG, cell, MC, and the UMI position together. cache_umi_position is deleted; aux_data_offset_from_record is now called exactly once per record on the decode hot path. Both improvements affect every command that consumes the UMI cache via the parallel pipeline (currently fgumi group from #343 and fgumi dedup from this commit). Mechanics: - AuxStringTags gains umi_position: Option<(u32, u16)> (aux-relative value offset + length, NUL excluded). - extract_aux_string_tags takes umi_tag: Option<[u8; 2]>; the inner loop captures UMI position in the same branch ladder as RG/cell/MC; early-exit mask widens from 7 to 15 when umi_tag is requested. - compute_group_key_from_raw returns (GroupKey, Option<(u32, u16)>); decode_records converts aux-relative to record-relative and stores via decoded.set_cached_umi(...). - Single-threaded group fast path (commands/group.rs:1408) passes None for umi_tag — it doesn't populate the cache and didn't before. Measurements (test workload, --threads 4, warm-cache tricord, profiling-profile binary): fgumi dedup (3-run median, runs within 0.02s): parent (PR343) this commit delta wall (s) 17.82 17.67 -0.8% cpu_time (s) 68.83 67.84 -1.4% max_rss (MB) 1435 1435 flat fgumi group: no detectable wall improvement (8-run min matches parent at 17.66s; per-record CPU saving below the noise floor at this workload size). Improvement remains architecturally real. Output BAM byte-identical pre/post on dedup across 53.1M records.
Resolves #334. Two improvements landed together because the smaller one revealed the larger one: 1. fgumi dedup now opts in to the per-record UMI position cache (#343 infrastructure) and reads via template.cached_umi() in its UMI- assignment loop with a fallback to find_string_tag. Skipped in --no-umi mode where the assignment site emits String::new(). 2. The cache populate, previously a separate second walk of aux data via cache_umi_position, is now folded into extract_aux_string_tags so a single linear walk captures RG, cell, MC, and the UMI position together. cache_umi_position is deleted; aux_data_offset_from_record is now called exactly once per record on the decode hot path. Both improvements affect every command that consumes the UMI cache via the parallel pipeline (currently fgumi group from #343 and fgumi dedup from this commit). Mechanics: - AuxStringTags gains umi_position: Option<(u32, u16)> (aux-relative value offset + length, NUL excluded). - extract_aux_string_tags takes umi_tag: Option<[u8; 2]>; the inner loop captures UMI position in the same branch ladder as RG/cell/MC; early-exit mask widens from 7 to 15 when umi_tag is requested. - compute_group_key_from_raw returns (GroupKey, Option<(u32, u16)>); decode_records converts aux-relative to record-relative and stores via decoded.set_cached_umi(...). - Single-threaded group fast path (commands/group.rs:1408) passes None for umi_tag — it doesn't populate the cache and didn't before. Measurements (test workload, --threads 4, warm-cache tricord, profiling-profile binary): fgumi dedup (3-run median, runs within 0.02s): parent (PR343) this commit delta wall (s) 17.82 17.67 -0.8% cpu_time (s) 68.83 67.84 -1.4% max_rss (MB) 1435 1435 flat fgumi group: no detectable wall improvement (8-run min matches parent at 17.66s; per-record CPU saving below the noise floor at this workload size). Improvement remains architecturally real. Output BAM byte-identical pre/post on dedup across 53.1M records.
…#351) Resolves #334. Two improvements landed together because the smaller one revealed the larger one: 1. fgumi dedup now opts in to the per-record UMI position cache (#343 infrastructure) and reads via template.cached_umi() in its UMI- assignment loop with a fallback to find_string_tag. Skipped in --no-umi mode where the assignment site emits String::new(). 2. The cache populate, previously a separate second walk of aux data via cache_umi_position, is now folded into extract_aux_string_tags so a single linear walk captures RG, cell, MC, and the UMI position together. cache_umi_position is deleted; aux_data_offset_from_record is now called exactly once per record on the decode hot path. Both improvements affect every command that consumes the UMI cache via the parallel pipeline (currently fgumi group from #343 and fgumi dedup from this commit). Mechanics: - AuxStringTags gains umi_position: Option<(u32, u16)> (aux-relative value offset + length, NUL excluded). - extract_aux_string_tags takes umi_tag: Option<[u8; 2]>; the inner loop captures UMI position in the same branch ladder as RG/cell/MC; early-exit mask widens from 7 to 15 when umi_tag is requested. - compute_group_key_from_raw returns (GroupKey, Option<(u32, u16)>); decode_records converts aux-relative to record-relative and stores via decoded.set_cached_umi(...). - Single-threaded group fast path (commands/group.rs:1408) passes None for umi_tag — it doesn't populate the cache and didn't before. Measurements (test workload, --threads 4, warm-cache tricord, profiling-profile binary): fgumi dedup (3-run median, runs within 0.02s): parent (PR343) this commit delta wall (s) 17.82 17.67 -0.8% cpu_time (s) 68.83 67.84 -1.4% max_rss (MB) 1435 1435 flat fgumi group: no detectable wall improvement (8-run min matches parent at 17.66s; per-record CPU saving below the noise floor at this workload size). Improvement remains architecturally real. Output BAM byte-identical pre/post on dedup across 53.1M records.
…plicate aux scan (#334/#343) (#405) The parallel Decode step already scans each record's aux data for the group key (RG/CB/MC). After this change, when a Group stage is in the chain it also records the UMI tag's value position on each DecodedRecord (`umi_value_offset: u32`, `umi_value_len: u16`). Template construction propagates the cached position to `Template::cached_umi_position`, and the UMI- assignment step in `fgumi group` slices the value through `Template::cached_umi` instead of calling `find_string_tag` a second time per template in the hot assignment loop. Wiring onto the #330 pipeline (the decode path moved into fgumi-bam-io + pipeline/steps/parse since #334/#343 were authored, so this is a re-application, not a cherry-pick): - fgumi-bam-io: DecodedRecord gains the cached UMI offset/len fields, the UMI_OFFSET_UNCACHED sentinel, and set_cached_umi/cached_umi/ cached_umi_position; GroupKeyConfig gains umi_tag + with_umi_tag. - parse/decode.rs: both decode paths (DecodeRecords, DecodeFromRecords) call a new cache_umi_position (using the already-present fgumi_raw_bam::find_string_tag_position) when umi_tag is set. - template.rs: Template::from_decoded_records + cached_umi pick the cached position from the primary read (R1, else R2; supplementary/secondary skipped); grouper::build_templates_from_records routes DecodedRecords through it. - chains/builder.rs: bam_group_key_config enables RX caching only when a Group stage is present, so non-group chains skip the extra per-record scan. The cache is independent of #330's name-hash-only key optimization (group key vs UMI value). Records built outside the Decode path (tests, SAM input, the single-threaded group path) leave the cache unset and fall back to find_string_tag, so output is byte-identical: a fallback-agreement unit test asserts cached_umi == find_string_tag, and the group/dedup determinism and runall parity suites pass unchanged.
…plicate aux scan (#334/#343) (#405) The parallel Decode step already scans each record's aux data for the group key (RG/CB/MC). After this change, when a Group stage is in the chain it also records the UMI tag's value position on each DecodedRecord (`umi_value_offset: u32`, `umi_value_len: u16`). Template construction propagates the cached position to `Template::cached_umi_position`, and the UMI- assignment step in `fgumi group` slices the value through `Template::cached_umi` instead of calling `find_string_tag` a second time per template in the hot assignment loop. Wiring onto the #330 pipeline (the decode path moved into fgumi-bam-io + pipeline/steps/parse since #334/#343 were authored, so this is a re-application, not a cherry-pick): - fgumi-bam-io: DecodedRecord gains the cached UMI offset/len fields, the UMI_OFFSET_UNCACHED sentinel, and set_cached_umi/cached_umi/ cached_umi_position; GroupKeyConfig gains umi_tag + with_umi_tag. - parse/decode.rs: both decode paths (DecodeRecords, DecodeFromRecords) call a new cache_umi_position (using the already-present fgumi_raw_bam::find_string_tag_position) when umi_tag is set. - template.rs: Template::from_decoded_records + cached_umi pick the cached position from the primary read (R1, else R2; supplementary/secondary skipped); grouper::build_templates_from_records routes DecodedRecords through it. - chains/builder.rs: bam_group_key_config enables RX caching only when a Group stage is present, so non-group chains skip the extra per-record scan. The cache is independent of #330's name-hash-only key optimization (group key vs UMI value). Records built outside the Decode path (tests, SAM input, the single-threaded group path) leave the cache unset and fall back to find_string_tag, so output is byte-identical: a fallback-agreement unit test asserts cached_umi == find_string_tag, and the group/dedup determinism and runall parity suites pass unchanged.
Summary
Resolves #334.
fgumi_raw_bam::find_tag_positionaccounts for ~5% of CPU time onfgumi groupbecause the same record's aux data is scanned independently by the Decode step (RG/CB/MC), the filter step (MQ/UMI), and the UMI-assignment step (UMI again). This change caches the UMI tag's value position on eachDecodedRecordduring the existing parallel Decode pass and slices the value out of that cache in the assignment step, eliminating the secondfind_tag_positionwalk per template.fgumi_raw_bam::find_string_tag_position(aux, tag) -> Option<(u32, u16)>returns the offset (within aux data) of a Z-tag's value bytes and its length, without the trailing NUL — useful to anything that wants to cache a Z-tag position rather than borrow its slice.GroupKeyConfiggains an opt-inumi_tag: Option<[u8; 2]>and awith_umi_tagbuilder.DecodedRecordcarriesumi_value_offset: u32(sentinelUMI_OFFSET_UNCACHED = u32::MAX) andumi_value_len: u16, populated bydecode_recordswhenumi_tagis set.Template::from_decoded_recordspicks the cached UMI from the primary read (R1, falling back to R2; secondary / supplementary records ignored) and stores it incached_umi_position.Template::cached_umi()slices the primary record bytes to recover the UMI without rescanning.Template::from_records(used by tests and non-group callers) leaves the cacheNone, soassign_umi_groups_for_indices_implfalls back tofind_string_tagexactly as before.fgumi groupopts in viaGroupKeyConfig::new(..).with_umi_tag(raw_tag).Test plan
cargo ci-fmtcargo ci-lintcargo ci-test— 2078 tests pass, including the existingtest_paired_mi_deterministicmatrix that exercisesfgumi groupend-to-end with multiple thread configurationstags::tests::test_find_string_tag_position_*(6 cases: first/subsequent tag, empty value, absent, non-Z, missing NUL)unified_pipeline::base::tests::test_decoded_record_cached_umi_{default_absent,round_trip,out_of_range_returns_none}unified_pipeline::bam::tests::test_cache_umi_position_{records_value_offset_and_len,leaves_cache_unset_when_tag_missing}andtest_decode_records_{populates,leaves}_umi_cache_*template::tests::test_template_from_decoded_records_{propagates_r1_umi_cache,falls_back_to_r2_when_no_r1,no_cache_when_primary_uncached,ignores_supplementary}target/profiling/fgumi groupwith samply at 999 Hz before/after onagilent-hs2.mapped.bamper the issue's verification plan; expectfind_tag_positioncontribution inassign_umi_groups_for_indices_implto drop to near zero (Decode-stage scan unchanged).