#![allow(dead_code)]
use crate::sparse_io::*;
use indicatif::{ParallelProgressIterator, ProgressIterator};
use legume_numeric::matrix::knn_match::ColumnDict;
use legume_numeric::matrix::knn_match::MakeVecPoint;
use legume_numeric::matrix::traits::*;
use legume_numeric::matrix::utils::*;
use log::info;
use rayon::prelude::*;
use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
use std::borrow::Cow;
use std::ops::Index;
use std::sync::Arc;
mod batch;
mod groups;
mod matched;
mod push;
mod read;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BackendLocation {
pub backend: u32,
pub local_col: u32,
}
type SparseData = dyn SparseIo<IndexIter = Vec<usize>>;
pub type RowNameCanonicalizer = Arc<dyn Fn(&str) -> Box<str> + Send + Sync>;
pub type ColumnNameCanonicalizer = Arc<dyn Fn(&str) -> Box<str> + Send + Sync>;
#[derive(Default)]
struct DerivedCaches {
col_to_group: Option<HashMap<usize, usize>>,
group_to_cols: Option<Vec<Vec<usize>>>,
group_keys: Option<Vec<Box<str>>>,
batch_knn_lookup: Option<Vec<ColumnDict<usize>>>,
col_to_batch: Option<Vec<usize>>,
batch_to_cols: Option<Vec<Vec<usize>>>,
batch_idx_to_name: Option<Vec<Box<str>>>,
between_batch_proximity: Option<Vec<Vec<usize>>>,
col_multiplicity: Option<Vec<f32>>,
}
impl Clone for DerivedCaches {
fn clone(&self) -> Self {
Self::default()
}
}
#[derive(Clone)]
pub struct SparseIoVec {
data_vec: Vec<Arc<SparseData>>,
col_to_data: Vec<Vec<BackendLocation>>,
data_to_cols: HashMap<usize, Vec<usize>>,
offset: usize,
row_canonicalizer: Option<RowNameCanonicalizer>,
row_name_position: HashMap<Box<str>, usize>,
row_names_by_global: Vec<Box<str>>,
data_local_to_global_row: Vec<Vec<usize>>,
data_global_to_local_row: Vec<HashMap<usize, usize>>,
data_has_intra_row_merges: Vec<bool>,
row_count_by_global: Vec<usize>,
global_to_compact_row: Vec<Option<usize>>,
compact_to_global_row: Vec<usize>,
column_names_with_data_tag: Vec<Box<str>>,
col_name_position: HashMap<Box<str>, u32>,
derived: DerivedCaches,
cached_num_rows: usize,
cached_num_columns: usize,
row_alignment: RowAlignment,
column_alignment: ColumnAlignment,
column_canonicalizer: Option<ColumnNameCanonicalizer>,
per_backend_row_suffix: Option<Vec<Box<str>>>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RowAlignment {
#[default]
Union,
Intersect,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum ColumnAlignment {
#[default]
Disjoint,
Union,
}
pub struct TripletsMatched {
pub shape: (usize, usize),
pub triplets: Vec<(u64, u64, f32)>,
pub source_columns: Vec<usize>,
pub matched_columns: Vec<usize>,
pub distances: Vec<f32>,
}
impl Index<usize> for SparseIoVec {
type Output = Arc<SparseData>;
fn index(&self, idx: usize) -> &Self::Output {
&self.data_vec[idx]
}
}
impl Default for SparseIoVec {
fn default() -> Self {
Self::new()
}
}
impl SparseIoVec {
pub fn new() -> Self {
Self {
data_vec: vec![],
col_to_data: vec![],
data_to_cols: HashMap::default(),
offset: 0,
row_canonicalizer: None,
row_name_position: HashMap::default(),
row_names_by_global: vec![],
data_local_to_global_row: vec![],
data_global_to_local_row: vec![],
data_has_intra_row_merges: vec![],
row_count_by_global: vec![],
global_to_compact_row: vec![],
compact_to_global_row: vec![],
column_names_with_data_tag: vec![],
col_name_position: HashMap::default(),
derived: DerivedCaches::default(),
cached_num_rows: 0,
cached_num_columns: 0,
row_alignment: RowAlignment::default(),
column_alignment: ColumnAlignment::default(),
column_canonicalizer: None,
per_backend_row_suffix: None,
}
}
pub fn with_row_alignment(mut self, mode: RowAlignment) -> anyhow::Result<Self> {
anyhow::ensure!(
self.data_vec.is_empty(),
"row alignment must be set before any push"
);
self.row_alignment = mode;
Ok(self)
}
pub fn with_row_canonicalizer(
mut self,
canon: impl Fn(&str) -> Box<str> + Send + Sync + 'static,
) -> anyhow::Result<Self> {
anyhow::ensure!(
self.data_vec.is_empty(),
"row canonicalizer must be set before any push"
);
self.row_canonicalizer = Some(Arc::new(canon));
Ok(self)
}
pub fn with_per_backend_row_suffix(mut self, suffix: Vec<Box<str>>) -> anyhow::Result<Self> {
anyhow::ensure!(
self.data_vec.is_empty(),
"per-backend row suffix must be set before any push"
);
self.per_backend_row_suffix = Some(suffix);
Ok(self)
}
pub fn with_column_alignment(mut self, mode: ColumnAlignment) -> anyhow::Result<Self> {
anyhow::ensure!(
self.data_vec.is_empty(),
"column alignment must be set before any push"
);
self.column_alignment = mode;
Ok(self)
}
pub fn with_column_canonicalizer(
mut self,
canon: impl Fn(&str) -> Box<str> + Send + Sync + 'static,
) -> anyhow::Result<Self> {
anyhow::ensure!(
self.data_vec.is_empty(),
"column canonicalizer must be set before any push"
);
self.column_canonicalizer = Some(Arc::new(canon));
Ok(self)
}
pub fn len(&self) -> usize {
self.data_vec.len()
}
pub fn is_empty(&self) -> bool {
self.data_vec.is_empty()
}
pub fn num_rows(&self) -> usize {
self.cached_num_rows
}
pub fn num_rows_in_at_least(&self, k: usize) -> usize {
self.row_count_by_global.iter().filter(|&&c| c >= k).count()
}
pub fn column_alignment(&self) -> ColumnAlignment {
self.column_alignment
}
#[must_use]
pub fn row_coverage_by_backend(&self) -> Option<Vec<Vec<bool>>> {
let n = self.cached_num_rows;
let mut coverage = vec![vec![false; n]; self.data_vec.len()];
let mut any_gap = false;
for (d, locals) in self.data_local_to_global_row.iter().enumerate() {
for &g in locals {
if let Some(r) = self.global_to_compact_row[g] {
coverage[d][r] = true;
}
}
any_gap |= coverage[d].iter().any(|&c| !c);
}
any_gap.then_some(coverage)
}
#[must_use]
pub fn column_source(&self, col: usize) -> Option<usize> {
match self.col_to_data.get(col)?.as_slice() {
[one] => Some(one.backend as usize),
_ => None,
}
}
#[must_use]
pub fn column_locations(&self, col: usize) -> &[BackendLocation] {
self.col_to_data.get(col).map_or(&[], Vec::as_slice)
}
pub fn num_non_zeros(&self) -> anyhow::Result<usize> {
let mut ret = 0;
for dat in self.data_vec.iter() {
let nnz = dat
.num_non_zeros()
.ok_or(anyhow::anyhow!("can't figure out the number of non-zeros"))?;
ret += nnz;
}
Ok(ret)
}
pub fn num_columns(&self) -> usize {
self.cached_num_columns
}
pub fn clone_for_collapse(&self) -> Self {
self.clone()
}
pub fn mask_rows(&mut self, keep: &[bool]) -> anyhow::Result<()> {
let n_compact = self.cached_num_rows;
if keep.len() != n_compact {
return Err(anyhow::anyhow!(
"mask_rows: keep.len()={} != num_rows={}",
keep.len(),
n_compact
));
}
let mut old_to_new: Vec<Option<usize>> = vec![None; n_compact];
let mut next = 0usize;
for (old, &k) in keep.iter().enumerate() {
if k {
old_to_new[old] = Some(next);
next += 1;
}
}
for entry in self.global_to_compact_row.iter_mut() {
*entry = entry.and_then(|old_compact| old_to_new[old_compact]);
}
self.cached_num_rows = next;
self.compact_to_global_row.clear();
self.compact_to_global_row.resize(next, 0);
for (g, &c_opt) in self.global_to_compact_row.iter().enumerate() {
if let Some(c) = c_opt {
self.compact_to_global_row[c] = g;
}
}
log::info!(
"mask_rows: {} → {} rows ({} excluded)",
n_compact,
next,
n_compact - next
);
Ok(())
}
pub fn mask_columns(&mut self, keep: &[bool]) -> anyhow::Result<()> {
let n = self.cached_num_columns;
if keep.len() != n {
return Err(anyhow::anyhow!(
"mask_columns: keep.len()={} != num_columns={}",
keep.len(),
n
));
}
debug_assert!(
self.derived.col_to_batch.is_none() && self.derived.col_to_group.is_none(),
"mask_columns must be called before batch/group registration"
);
let mut old_to_new: Vec<Option<usize>> = vec![None; n];
let mut next = 0usize;
for (old, &k) in keep.iter().enumerate() {
if k {
old_to_new[old] = Some(next);
next += 1;
}
}
let mut new_col_to_data: Vec<Vec<BackendLocation>> = Vec::with_capacity(next);
let mut new_names: Vec<Box<str>> = Vec::with_capacity(next);
for (old, &k) in keep.iter().enumerate() {
if k {
new_col_to_data.push(std::mem::take(&mut self.col_to_data[old]));
new_names.push(self.column_names_with_data_tag[old].clone());
}
}
self.col_to_data = new_col_to_data;
self.column_names_with_data_tag = new_names;
for cols in self.data_to_cols.values_mut() {
for c in cols.iter_mut() {
*c = if *c == usize::MAX {
usize::MAX
} else {
old_to_new[*c].unwrap_or(usize::MAX)
};
}
}
if !self.col_name_position.is_empty() {
let mut pos: HashMap<Box<str>, u32> =
HashMap::with_capacity_and_hasher(next, Default::default());
for (g, name) in self.column_names_with_data_tag.iter().enumerate() {
let canon: Box<str> = match self.column_canonicalizer.as_ref() {
Some(c) => c(name),
None => name.clone(),
};
let g_u32: u32 = g
.try_into()
.map_err(|_| anyhow::anyhow!("global col overflows u32"))?;
pos.insert(canon, g_u32);
}
self.col_name_position = pos;
}
self.offset = next;
self.cached_num_columns = next;
self.derived = DerivedCaches::default();
log::info!(
"mask_columns: {} → {} cells ({} excluded)",
n,
next,
n - next
);
Ok(())
}
pub fn clear_column_membership(&mut self) {
self.derived = DerivedCaches::default();
}
pub fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
let ntot = self.num_rows();
let mut ret = vec![Box::from(""); ntot];
for (raw_global, name) in self.row_names_by_global.iter().enumerate() {
if let Some(compact) = self
.global_to_compact_row
.get(raw_global)
.copied()
.flatten()
{
ret[compact] = name.clone();
}
}
Ok(ret)
}
}