vcf_reformatter/
extract_sample_info.rs

1use std::collections::HashMap;
2
3#[derive(Debug, Clone)]
4pub struct ParsedSample {
5    pub sample_name: String,
6    pub format_fields: HashMap<String, String>,
7}
8
9#[derive(Debug, Clone)]
10pub struct ParsedFormatSample {
11    pub format_keys: Vec<String>,
12    pub samples: Vec<ParsedSample>,
13}
14
15impl Default for ParsedFormatSample {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl ParsedFormatSample {
22    pub fn new() -> Self {
23        ParsedFormatSample {
24            format_keys: Vec::new(),
25            samples: Vec::new(),
26        }
27    }
28    pub fn from_vcf_fields(
29        fields: &[&str],
30        column_names: &[&str],
31    ) -> Result<Self, Box<dyn std::error::Error>> {
32        if fields.is_empty() {
33            return Ok(Self::new());
34        }
35
36        let format_field = fields.first().copied();
37        let sample_fields: Vec<String> = fields.iter().skip(1).map(|s| s.to_string()).collect();
38
39        // Get sample names from column names, starting after the 9 standard VCF columns
40        let sample_names: Vec<String> = if column_names.len() > 9 {
41            column_names[9..].iter().map(|s| s.to_string()).collect()
42        } else {
43            // Generate default sample names if not provided
44            (0..sample_fields.len())
45                .map(|i| format!("SAMPLE_{}", i + 1))
46                .collect()
47        };
48
49        parse_format_and_samples(format_field, &sample_fields, &sample_names)
50    }
51
52    pub fn get_headers_for_samples(&self) -> Vec<String> {
53        let mut headers = Vec::new();
54
55        for sample in &self.samples {
56            for format_key in &self.format_keys {
57                headers.push(format!("{}_{}", sample.sample_name, format_key));
58            }
59        }
60
61        headers
62    }
63}
64
65pub fn parse_format_and_samples(
66    format_field: Option<&str>,
67    sample_fields: &[String],
68    sample_names: &[String],
69) -> Result<ParsedFormatSample, Box<dyn std::error::Error>> {
70    let mut parsed = ParsedFormatSample::new();
71
72    if let Some(format_str) = format_field {
73        // Skip if format string is empty or just whitespace
74        if format_str.trim().is_empty() {
75            return Ok(parsed);
76        }
77
78        parsed.format_keys = format_str.split(':').map(|s| s.to_string()).collect();
79
80        for (i, sample_field) in sample_fields.iter().enumerate() {
81            let sample_name = sample_names
82                .get(i)
83                .map(|s| s.to_string())
84                .unwrap_or_else(|| format!("SAMPLE_{}", i + 1));
85
86            let sample_values: Vec<&str> = sample_field.split(':').collect();
87            let mut format_fields = HashMap::new();
88
89            for (j, format_key) in parsed.format_keys.iter().enumerate() {
90                let value = sample_values.get(j).unwrap_or(&".").to_string();
91                format_fields.insert(format_key.clone(), value);
92            }
93
94            parsed.samples.push(ParsedSample {
95                sample_name,
96                format_fields,
97            });
98        }
99    }
100
101    Ok(parsed)
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn test_parse_format_and_samples() {
110        let format = Some("GT:DP:AD:RO:QR:AO:QA:GL");
111        let sample_fields = vec!["1/1:19:0,19:0:0:19:579:-32.2782,-5.71957,0".to_string()];
112        let sample_names = vec!["B505_B505_V_1".to_string()];
113
114        let parsed = parse_format_and_samples(format, &sample_fields, &sample_names).unwrap();
115
116        assert_eq!(
117            parsed.format_keys,
118            vec!["GT", "DP", "AD", "RO", "QR", "AO", "QA", "GL"]
119        );
120        assert_eq!(parsed.samples.len(), 1);
121        assert_eq!(parsed.samples[0].sample_name, "B505_B505_V_1");
122        assert_eq!(parsed.samples[0].format_fields.get("GT").unwrap(), "1/1");
123        assert_eq!(parsed.samples[0].format_fields.get("DP").unwrap(), "19");
124        assert_eq!(parsed.samples[0].format_fields.get("AD").unwrap(), "0,19");
125    }
126
127    #[test]
128    fn test_get_headers_for_samples() {
129        let mut parsed = ParsedFormatSample::new();
130        parsed.format_keys = vec!["GT".to_string(), "DP".to_string()];
131
132        let sample = |name: &str, gt: &str, dp: &str| ParsedSample {
133            sample_name: name.to_string(),
134            format_fields: HashMap::from([
135                ("GT".to_string(), gt.to_string()),
136                ("DP".to_string(), dp.to_string()),
137            ]),
138        };
139        parsed.samples = vec![
140            sample("SAMPLE1", "0/1", "20"),
141            sample("SAMPLE2", "1/1", "30"),
142        ];
143
144        let headers = parsed.get_headers_for_samples();
145        assert_eq!(
146            headers,
147            vec!["SAMPLE1_GT", "SAMPLE1_DP", "SAMPLE2_GT", "SAMPLE2_DP"]
148        );
149    }
150}
vcf_reformatter/extract_sample_info.rs

vcf_reformatter/
extract_sample_info.rs
