annonars 0.26.0

Rust template repository
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

//! Import of gonosomal variant frequencies.

use crate::{common, freqs};

/// Write out the given record to the database.
fn write_record(
    db: &rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
    cf: &std::sync::Arc<rocksdb::BoundColumnFamily>,
    record_key: &common::keys::Var,
    record_genome: &mut Option<noodles_vcf::Record>,
    record_exome: &mut Option<noodles_vcf::Record>,
) -> Result<(), anyhow::Error> {
    if record_genome.is_none() && record_exome.is_none() {
        // Early exit, nothing to write out.
        return Ok(());
    }

    let count_genomes = if let Some(record_genome) = record_genome {
        freqs::serialized::xy::Counts::from_vcf_allele(record_genome, 0)
    } else {
        freqs::serialized::xy::Counts::default()
    };
    let counts_exomes = if let Some(record_exome) = record_exome {
        freqs::serialized::xy::Counts::from_vcf_allele(record_exome, 0)
    } else {
        freqs::serialized::xy::Counts::default()
    };

    let gono_record = freqs::serialized::xy::Record {
        gnomad_genomes: count_genomes,
        gnomad_exomes: counts_exomes,
    };

    let mut buf = vec![0u8; freqs::serialized::xy::Record::buf_len()];
    gono_record.to_buf(&mut buf);
    let key: Vec<u8> = record_key.clone().into();

    // tracing::info!("  key = {:?}, record = {:?}", &record_key, &gono_record);

    db.put_cf(cf, key, &buf)?;

    Ok(())
}

/// Import of gonosomal variant frequencies.
pub fn import_region(
    db: &rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
    path_genome: Option<&String>,
    path_exome: Option<&String>,
    region: &noodles_core::region::Region,
) -> Result<(), anyhow::Error> {
    // Get handle to "gonosomal" column family.
    let cf_gono = db.cf_handle("gonosomal").unwrap();
    // Build `Vec` of readers and by-index map that tells whether it is genomes.
    let mut is_genome = Vec::new();
    let mut readers = Vec::new();
    if let Some(path_genome) = path_genome {
        is_genome.push(true);
        readers.push(noodles_vcf::indexed_reader::Builder::default().build_from_path(path_genome)?);
    }
    if let Some(path_exome) = path_exome {
        is_genome.push(false);
        readers.push(noodles_vcf::indexed_reader::Builder::default().build_from_path(path_exome)?);
    }
    // Read headers.
    let headers: Vec<_> = readers
        .iter_mut()
        .map(|reader| reader.read_header())
        .collect::<Result<_, _>>()?;

    // Seek to region obtaining `ueries`.
    let queries: Vec<_> = readers
        .iter_mut()
        .zip(&headers)
        .map(|(reader, header)| reader.query(header, region))
        .collect::<Result<_, _>>()?;
    // Construct the `MultiQuery`.
    let multi_query = super::reading::MultiQuery::new(queries)?;

    // Now iterate over the `MultiQuery` and write to the database.
    //
    // The key of the records stored in `record_{genome,exome}`.
    let mut record_key = None;
    // Record from gnomAD genomes (same position as record_exome, if either).
    let mut record_genome = None;
    // Record from gnomAD exomes (same position as record_genome, if either).
    let mut record_exome = None;
    for result in multi_query {
        let (idx, record) = result?;
        // Obtain the key of the next record.
        let curr_key = common::keys::Var::from_vcf_allele(&record, 0);

        // Write out current records to database if we advance.
        if record_key.as_ref() != Some(&curr_key) {
            if let Some(record_key) = record_key.as_ref() {
                write_record(
                    db,
                    &cf_gono,
                    record_key,
                    &mut record_genome,
                    &mut record_exome,
                )?;
            }
            record_genome = None;
            record_exome = None;
        }

        // Shortcut to whether this from gnomAD genomes.
        if is_genome[idx] {
            record_genome = Some(record);
        } else {
            record_exome = Some(record);
        }

        // Updating the current key cannot cause harm.
        record_key = Some(curr_key);
    }

    // Write final records to database.
    if let Some(record_key) = record_key.as_ref() {
        write_record(
            db,
            &cf_gono,
            record_key,
            &mut record_genome,
            &mut record_exome,
        )?;
    }

    Ok(())
}