Skip to main content

data_beans/aux/
feature_types.rs

1//! The typed feature table that rides beside a mixed-type embedding:
2//! `{prefix}.feature_types.parquet`, string columns `feature` and `type`, one
3//! row per embedding row in the same order. `senna fne` and `gene-text` write
4//! it; a consumer that wants only one type of row (`senna bge` pinning gene
5//! rows, say) reads it. The type names are the shared vocabulary.
6
7use legume_numeric::matrix::parquet::{
8    read_parquet_string_columns_by_name, write_named_table, Column,
9};
10use std::path::Path;
11
12/// Nodes whose names are canonicalised as gene symbols.
13pub const GENE_TYPE: &str = "gene";
14/// Ontology terms and gene sets.
15pub const TERM_TYPE: &str = "term";
16/// Fixed genomic windows.
17pub const REGION_TYPE: &str = "region";
18/// Vocabulary words of a text relation.
19pub const WORD_TYPE: &str = "word";
20
21pub fn feature_types_path(prefix: &str) -> String {
22    format!("{prefix}.feature_types.parquet")
23}
24
25/// Write the table for `names[i]` of type `types[i]`.
26pub fn write_feature_types(
27    prefix: &str,
28    names: &[Box<str>],
29    types: &[Box<str>],
30) -> anyhow::Result<()> {
31    anyhow::ensure!(
32        names.len() == types.len(),
33        "feature types: {} names for {} types",
34        names.len(),
35        types.len()
36    );
37    write_named_table(
38        &feature_types_path(prefix),
39        "feature",
40        names,
41        &[(Box::from("type"), Column::Str(types))],
42    )
43}
44
45/// One row of the table: the feature's name and its type.
46pub type FeatureType = (Box<str>, Box<str>);
47
48/// The run's rows; `None` when the run wrote no table, which a caller reads as
49/// "every row is of the one type it expects".
50pub fn read_feature_types(prefix: &str) -> anyhow::Result<Option<Vec<FeatureType>>> {
51    let path = feature_types_path(prefix);
52    if !Path::new(&path).exists() {
53        return Ok(None);
54    }
55    let mut cols = read_parquet_string_columns_by_name(&path, &["feature", "type"])?;
56    let types = cols.pop().expect("two columns requested");
57    let names = cols.pop().expect("two columns requested");
58    Ok(Some(names.into_iter().zip(types).collect()))
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn round_trip_and_absence() {
67        let dir = tempfile::tempdir().unwrap();
68        let prefix = dir.path().join("run").to_string_lossy().into_owned();
69        assert!(read_feature_types(&prefix).unwrap().is_none());
70        let names: Vec<Box<str>> = vec!["TP53".into(), "GO:1".into()];
71        let types: Vec<Box<str>> = vec![GENE_TYPE.into(), TERM_TYPE.into()];
72        write_feature_types(&prefix, &names, &types).unwrap();
73        let rows = read_feature_types(&prefix).unwrap().unwrap();
74        assert_eq!(
75            rows,
76            vec![
77                ("TP53".into(), GENE_TYPE.into()),
78                ("GO:1".into(), TERM_TYPE.into())
79            ]
80        );
81        assert!(write_feature_types(&prefix, &names, &types[..1]).is_err());
82    }
83}