Frozen dictionaries: match only the source rows a caller marks

src/aux/feature_types.rs

@@ -18,6 +18,13 @@ pub const REGION_TYPE: &str = "region";
 /// Vocabulary words of a text relation.
 pub const WORD_TYPE: &str = "word";
 
+/// Whether rows of type `t` name a data feature — a gene, or a genomic window
+/// a peak can match — rather than a term, word or cell type.
+#[must_use]
+pub fn is_data_feature_type(t: &str) -> bool {
+    matches!(t, GENE_TYPE | REGION_TYPE)
+}
+
 pub fn feature_types_path(prefix: &str) -> String {
     format!("{prefix}.feature_types.parquet")
 }
@@ -42,6 +49,18 @@ pub fn write_feature_types(
     )
 }
 
+/// One flag per row of a table (`names`, in order): true for a gene or a
+/// genomic window, by the table's types table `types`. `None` when `types`
+/// does not list `names` in order — written for another table under the
+/// same prefix. For [`crate::aux::frozen_features::load_frozen_feature_host_matching`].
+#[must_use]
+pub fn feature_rows(types: &[FeatureType], names: &[Box<str>]) -> Option<Vec<bool>> {
+    if !types.iter().map(|(n, _)| n).eq(names.iter()) {
+        return None;
+    }
+    Some(types.iter().map(|(_, t)| is_data_feature_type(t)).collect())
+}
+
 /// One row of the table: the feature's name and its type.
 pub type FeatureType = (Box<str>, Box<str>);
 
@@ -62,6 +81,22 @@ pub fn read_feature_types(prefix: &str) -> anyhow::Result<Option<Vec<FeatureType
 mod tests {
     use super::*;
 
+    #[test]
+    fn feature_rows_are_marked_by_position_for_the_table_listed() {
+        let t = |n: &str, ty: &str| -> FeatureType { (n.into(), ty.into()) };
+        let types = [
+            t("CD4", "cell_type"),
+            t("CD4", "gene"),
+            t("chr1:0-5000", "region"),
+        ];
+        let names: Vec<Box<str>> = vec!["CD4".into(), "CD4".into(), "chr1:0-5000".into()];
+        assert_eq!(feature_rows(&types, &names), Some(vec![false, true, true]));
+        assert_eq!(feature_rows(&types, &names[..2]), None);
+        let other: Vec<Box<str>> = vec!["CD4".into(), "MYC".into(), "chr1:0-5000".into()];
+        assert_eq!(feature_rows(&types, &other), None);
+        assert!(is_data_feature_type(GENE_TYPE) && !is_data_feature_type(WORD_TYPE));
+    }
+
     #[test]
     fn round_trip_and_absence() {
         let dir = tempfile::tempdir().unwrap();

src/aux/frozen_features.rs

@@ -76,14 +76,33 @@ pub struct FrozenLoadArgs<'a> {
     /// matching; [`FeatureNameKind::Gene { delim: '_' }`] is the typical
     /// choice for scRNA gene IDs.
     pub name_kind: FeatureNameKind,
-    /// Applied to every SOURCE row name before canonicalization, and kept as
+    /// Applied to each SOURCE row name that may match (every row, unless
+    /// [`load_frozen_feature_host_matching`] marks fewer) before
+    /// canonicalization, and kept as
     /// the host's `src_names`: how a caller whose axis carries a row grammar
     /// (`{gene}/count/spliced`) reads a plain gene table, lifting each bare
     /// name into the grammar first. `None` = the names as read.
     pub source_name_map: Option<SourceNameMap<'a>>,
 }
 
+/// Load the dictionary and match its rows to the target axis by canonical
+/// name. Among source rows with one canonical name the first wins.
 pub fn load_frozen_feature_host(args: FrozenLoadArgs) -> anyhow::Result<FrozenFeatureHost> {
+    load_frozen_feature_host_matching(args, |names| Ok(vec![true; names.len()]))
+}
+
+/// [`load_frozen_feature_host`], matching only the source rows `matchable`
+/// marks. It is handed the dictionary's row names as read (before
+/// [`FrozenLoadArgs::source_name_map`]) and returns one flag per row, BY
+/// POSITION: two rows may share a name (a cell type `CD4` beside the gene
+/// `CD4`), so marking by name would mark both and the first would still win.
+/// [`crate::aux::feature_types::feature_rows`] marks a mixed-type table's
+/// gene and region rows from its types table. A row left unmarked is never
+/// matched or renamed, yet stays in `src_names` / `src_e_feat`.
+pub fn load_frozen_feature_host_matching(
+    args: FrozenLoadArgs,
+    matchable: impl FnOnce(&[Box<str>]) -> anyhow::Result<Vec<bool>>,
+) -> anyhow::Result<FrozenFeatureHost> {
     let dict = <DMatrix<f32> as IoOps>::from_parquet(args.dictionary_path)?;
     let n_src = dict.rows.len();
     let h = dict.mat.ncols();
@@ -116,15 +135,41 @@ pub fn load_frozen_feature_host(args: FrozenLoadArgs) -> anyhow::Result<FrozenFe
         }
     };
 
+    let matchable = matchable(&dict.rows)
+        .map_err(|e| anyhow::anyhow!("{}: marking its rows: {e}", args.dictionary_path))?;
+    anyhow::ensure!(
+        matchable.len() == n_src,
+        "{}: {} row flags for {} rows",
+        args.dictionary_path,
+        matchable.len(),
+        n_src
+    );
+    let n_matchable = matchable.iter().filter(|&&m| m).count();
+    anyhow::ensure!(
+        n_src == 0 || n_matchable > 0,
+        "{}: none of its {} rows is marked as one that may match a feature",
+        args.dictionary_path,
+        n_src
+    );
+    // Only a row that may match is renamed: an unmarked one keeps its name.
     let src_names: Vec<Box<str>> = match args.source_name_map {
-        Some(f) => dict.rows.iter().map(|n| f(n)).collect(),
+        Some(f) => dict
+            .rows
+            .iter()
+            .zip(&matchable)
+            .map(|(n, &m)| if m { f(n) } else { n.clone() })
+            .collect(),
         None => dict.rows,
     };
     let mut src_by_canon: FxHashMap<Box<str>, usize> = FxHashMap::default();
     let mut src_dupes = 0usize;
     for (i, name) in src_names.iter().enumerate() {
+        if !matchable[i] {
+            continue;
+        }
         let canon = args.name_kind.canonicalize(name);
-        // First occurrence wins, as documented; `insert` would keep the last.
+        // First occurrence wins (see `load_frozen_feature_host`); `insert`
+        // would keep the last.
         if let std::collections::hash_map::Entry::Vacant(e) = src_by_canon.entry(canon) {
             e.insert(i);
         } else {
@@ -138,6 +183,18 @@ pub fn load_frozen_feature_host(args: FrozenLoadArgs) -> anyhow::Result<FrozenFe
             src_dupes
         );
     }
+    let shadowing = src_names
+        .iter()
+        .zip(&matchable)
+        .filter(|(n, &m)| !m && src_by_canon.contains_key(&args.name_kind.canonicalize(n)))
+        .count();
+    if shadowing > 0 {
+        log::info!(
+            "{}: {} unmarked rows share a name with a matchable row and were passed over",
+            args.dictionary_path,
+            shadowing
+        );
+    }
 
     let mut keep_target_indices = Vec::new();
     let mut keep_src_indices = Vec::new();
@@ -167,24 +224,32 @@ pub fn load_frozen_feature_host(args: FrozenLoadArgs) -> anyhow::Result<FrozenFe
         .iter()
         .enumerate()
         .filter(|(i, r)| {
-            !unique_src_used.contains(i) && crate::aux::feature_rows::parse_feature_row(r).is_some()
+            matchable[*i]
+                && !unique_src_used.contains(i)
+                && crate::aux::feature_rows::parse_feature_row(r).is_some()
         })
         .count();
     if channelized_unmatched > 0 {
         log::warn!(
-            "{}: {} unmatched source rows carry the channelized row grammar —              is this a raw gene dictionary, or a channelized/co-embedding output?",
+            "{}: {} unmatched source rows carry the channelized row grammar — is this a raw gene dictionary, or a channelized/co-embedding output?",
             args.dictionary_path,
             channelized_unmatched
         );
     }
+    let matchable_note = if n_matchable < n_src {
+        format!("; {n_matchable} source rows may match")
+    } else {
+        String::new()
+    };
     log::info!(
-        "Frozen feature side from {}: {}/{} target features matched (H={}, {} of {} source rows reused, kind={:?})",
+        "Frozen feature side from {}: {}/{} target features matched (H={}, {} of {} source rows reused{}, kind={:?})",
         args.dictionary_path,
         keep_target_indices.len(),
         args.target_feature_names.len(),
         h,
         unique_src_used.len(),
         n_src,
+        matchable_note,
         args.name_kind
     );
 
@@ -365,6 +430,79 @@ mod tests {
         assert_eq!(host.e_feat[(1, 0)], 1.0);
     }
 
+    /// An unmarked row is never matched, though it comes first and shares
+    /// the gene's name, and stays in the source table.
+    #[test]
+    fn only_the_marked_rows_match() {
+        let dir = tempfile::tempdir().unwrap();
+        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
+        let src = DMatrix::<f32>::from_row_slice(3, 2, &[9.0, 9.0, 1.0, 2.0, 3.0, 4.0]);
+        write_test_parquet(
+            &dict_path,
+            &["CD4", "CD4", "MYC"],
+            "feature",
+            &["h0", "h1"],
+            &src,
+        );
+        let target: Vec<Box<str>> = ["MYC", "CD4"].iter().map(|s| (*s).into()).collect();
+        let args = || FrozenLoadArgs {
+            dictionary_path: &dict_path,
+            bias_path: None,
+            target_feature_names: &target,
+            name_kind: FeatureNameKind::Exact,
+            source_name_map: None,
+        };
+        let host =
+            load_frozen_feature_host_matching(args(), |_| Ok(vec![false, true, true])).unwrap();
+        assert_eq!(host.keep_target_indices, vec![0, 1]);
+        assert_eq!(host.keep_src_indices, vec![2, 1]);
+        assert_eq!(
+            host.e_feat.row(1).iter().copied().collect::<Vec<_>>(),
+            [1.0, 2.0]
+        );
+        assert_eq!(host.src_names.len(), 3);
+
+        // Every row unmarked, or the wrong count: refused.
+        let err = |m: Vec<bool>| {
+            load_frozen_feature_host_matching(args(), |_| Ok(m))
+                .err()
+                .unwrap()
+                .to_string()
+        };
+        assert!(err(vec![false; 3]).contains("none of its 3 rows is marked"));
+        assert!(err(vec![true]).contains("1 row flags for 3 rows"));
+
+        // The marks are asked of the names as read, before any rename.
+        let lift = |n: &str| -> Box<str> { format!("{n}/count/spliced").into() };
+        let mut seen: Vec<Box<str>> = Vec::new();
+        let renamed = FrozenLoadArgs {
+            source_name_map: Some(&lift),
+            ..args()
+        };
+        assert!(load_frozen_feature_host_matching(renamed, |names| {
+            seen = names.to_vec();
+            Ok(vec![false; names.len()])
+        })
+        .is_err());
+        let read: Vec<Box<str>> = vec!["CD4".into(), "CD4".into(), "MYC".into()];
+        assert_eq!(seen, read);
+
+        // ...and only a marked row is renamed.
+        let target: Vec<Box<str>> = vec!["CD4/count/spliced".into()];
+        let host = load_frozen_feature_host_matching(
+            FrozenLoadArgs {
+                target_feature_names: &target,
+                source_name_map: Some(&lift),
+                ..args()
+            },
+            |_| Ok(vec![false, true, true]),
+        )
+        .unwrap();
+        assert_eq!(&*host.src_names[0], "CD4");
+        assert_eq!(&*host.src_names[1], "CD4/count/spliced");
+        assert_eq!(host.keep_src_indices, vec![1]);
+    }
+
     #[test]
     fn empty_intersection_errors() {
         let dir = tempfile::tempdir().unwrap();