1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::alignment::{SampleAlignmentPlan, SourceSampleSet};
6use crate::error::{DataError, Result};
7use crate::handle::CoordinatorFeatureBlock;
8use crate::ids::{ObservationId, SampleId, SourceId};
9
10#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
11pub struct FeatureFusionPolicy {
12 #[serde(default = "default_true")]
13 pub namespace_columns: bool,
14}
15
16impl Default for FeatureFusionPolicy {
17 fn default() -> Self {
18 Self {
19 namespace_columns: true,
20 }
21 }
22}
23
24fn default_true() -> bool {
25 true
26}
27
28#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
29pub struct SourceFeatureBlock {
30 pub source_id: SourceId,
31 pub block: CoordinatorFeatureBlock,
32}
33
34pub fn source_sample_set_from_feature_block(block: &SourceFeatureBlock) -> Result<SourceSampleSet> {
35 validate_feature_block(&block.block)?;
36 let mut seen = BTreeSet::new();
37 let mut sample_ids = Vec::new();
38 for sample_id in &block.block.sample_ids {
39 if seen.insert(sample_id) {
40 sample_ids.push(sample_id.clone());
41 }
42 }
43 Ok(SourceSampleSet {
44 source_id: block.source_id.clone(),
45 sample_ids,
46 })
47}
48
49pub fn fuse_feature_blocks(
50 feature_set_id: impl Into<String>,
51 blocks: &[SourceFeatureBlock],
52 alignment: &SampleAlignmentPlan,
53 policy: &FeatureFusionPolicy,
54) -> Result<CoordinatorFeatureBlock> {
55 let feature_set_id = feature_set_id.into();
56 if feature_set_id.trim().is_empty() {
57 return Err(DataError::Validation(
58 "fused feature set id is empty".to_string(),
59 ));
60 }
61 if blocks.is_empty() {
62 return Err(DataError::Validation(
63 "feature fusion requires at least one source block".to_string(),
64 ));
65 }
66 alignment.validate()?;
67
68 let mut source_ids = BTreeSet::new();
69 for block in blocks {
70 if !source_ids.insert(&block.source_id) {
71 return Err(DataError::Validation(format!(
72 "feature fusion contains duplicate source `{}`",
73 block.source_id
74 )));
75 }
76 validate_feature_block(&block.block)?;
77 }
78 for mask in &alignment.masks {
79 if !source_ids.contains(&mask.source_id) {
80 return Err(DataError::Validation(format!(
81 "alignment mask references source `{}` absent from feature fusion",
82 mask.source_id
83 )));
84 }
85 }
86 if alignment.masks.len() != blocks.len() {
87 return Err(DataError::Validation(
88 "feature fusion sources and alignment masks differ".to_string(),
89 ));
90 }
91
92 let representation_id = blocks[0].block.representation_id.clone();
93 for block in blocks.iter().skip(1) {
94 if block.block.representation_id != representation_id {
95 return Err(DataError::Validation(format!(
96 "feature fusion source `{}` representation `{}` does not match `{}`",
97 block.source_id, block.block.representation_id, representation_id
98 )));
99 }
100 }
101
102 let mut feature_names = Vec::new();
103 for block in blocks {
104 for feature_name in &block.block.feature_names {
105 let output_name = if policy.namespace_columns {
106 format!("{}.{}", block.source_id, feature_name)
107 } else {
108 feature_name.clone()
109 };
110 feature_names.push(output_name);
111 }
112 }
113 let mut seen_features = BTreeSet::new();
114 for feature_name in &feature_names {
115 if !seen_features.insert(feature_name) {
116 return Err(DataError::Validation(format!(
117 "feature fusion produced duplicate feature `{feature_name}`"
118 )));
119 }
120 }
121
122 let row_maps = blocks
123 .iter()
124 .map(|block| (&block.source_id, rows_by_sample(&block.block)))
125 .collect::<BTreeMap<_, _>>();
126 validate_alignment_presence(blocks, alignment, &row_maps)?;
127 let reference = &blocks[0];
128 let reference_rows = row_maps
129 .get(&reference.source_id)
130 .expect("reference source map was created");
131
132 let mut observation_ids = Vec::new();
133 let mut sample_ids = Vec::new();
134 let mut values = Vec::new();
135
136 for sample_id in &alignment.sample_ids {
137 let output_rows = reference_rows
138 .get(sample_id)
139 .map(|indices| {
140 indices
141 .iter()
142 .map(|idx| OutputRow::Reference(*idx))
143 .collect::<Vec<_>>()
144 })
145 .unwrap_or_else(|| vec![OutputRow::Synthetic]);
146
147 for output_row in output_rows {
148 let mut row_values = Vec::new();
149 match output_row {
150 OutputRow::Reference(idx) => {
151 observation_ids.push(reference.block.observation_ids[idx].clone());
152 sample_ids.push(sample_id.clone());
153 }
154 OutputRow::Synthetic => {
155 observation_ids.push(synthetic_observation_id(sample_id)?);
156 sample_ids.push(sample_id.clone());
157 }
158 }
159 for block in blocks {
160 let source_rows = row_maps
161 .get(&block.source_id)
162 .expect("source row map was created");
163 if block.source_id == reference.source_id {
164 match output_row {
165 OutputRow::Reference(idx) => {
166 row_values.extend(block.block.values[idx].iter().cloned());
167 }
168 OutputRow::Synthetic => {
169 row_values.extend(std::iter::repeat_n(
170 serde_json::Value::Null,
171 block.block.feature_names.len(),
172 ));
173 }
174 }
175 continue;
176 }
177
178 match source_rows.get(sample_id).map(Vec::as_slice) {
179 Some([idx]) => row_values.extend(block.block.values[*idx].iter().cloned()),
180 Some(indices) => {
181 return Err(DataError::Validation(format!(
182 "feature fusion cannot broadcast {} repeated rows from non-reference source `{}` for sample `{sample_id}`",
183 indices.len(),
184 block.source_id
185 )));
186 }
187 None => row_values.extend(std::iter::repeat_n(
188 serde_json::Value::Null,
189 block.block.feature_names.len(),
190 )),
191 }
192 }
193 values.push(row_values);
194 }
195 }
196
197 let fused = CoordinatorFeatureBlock {
198 feature_set_id,
199 representation_id,
200 feature_names,
201 observation_ids,
202 sample_ids,
203 values,
204 };
205 validate_feature_block(&fused)?;
206 Ok(fused)
207}
208
209#[derive(Clone, Copy)]
210enum OutputRow {
211 Reference(usize),
212 Synthetic,
213}
214
215fn validate_feature_block(block: &CoordinatorFeatureBlock) -> Result<()> {
216 if block.feature_set_id.trim().is_empty() {
217 return Err(DataError::Validation(
218 "feature block feature_set_id is empty".to_string(),
219 ));
220 }
221 if block.feature_names.is_empty() {
222 return Err(DataError::Validation(format!(
223 "feature block `{}` contains no features",
224 block.feature_set_id
225 )));
226 }
227 if block.observation_ids.len() != block.sample_ids.len()
228 || block.sample_ids.len() != block.values.len()
229 {
230 return Err(DataError::Validation(format!(
231 "feature block `{}` row identity/value lengths differ",
232 block.feature_set_id
233 )));
234 }
235 let mut observations = BTreeSet::new();
236 for (idx, values) in block.values.iter().enumerate() {
237 if !observations.insert(&block.observation_ids[idx]) {
238 return Err(DataError::Validation(format!(
239 "feature block `{}` contains duplicate observation `{}`",
240 block.feature_set_id, block.observation_ids[idx]
241 )));
242 }
243 if values.len() != block.feature_names.len() {
244 return Err(DataError::Validation(format!(
245 "feature block `{}` row `{}` has {} values for {} features",
246 block.feature_set_id,
247 block.observation_ids[idx],
248 values.len(),
249 block.feature_names.len()
250 )));
251 }
252 }
253 Ok(())
254}
255
256fn rows_by_sample(block: &CoordinatorFeatureBlock) -> BTreeMap<&SampleId, Vec<usize>> {
257 let mut rows = BTreeMap::<&SampleId, Vec<usize>>::new();
258 for (idx, sample_id) in block.sample_ids.iter().enumerate() {
259 rows.entry(sample_id).or_default().push(idx);
260 }
261 rows
262}
263
264fn validate_alignment_presence<'a>(
265 blocks: &'a [SourceFeatureBlock],
266 alignment: &SampleAlignmentPlan,
267 row_maps: &BTreeMap<&'a SourceId, BTreeMap<&'a SampleId, Vec<usize>>>,
268) -> Result<()> {
269 for (idx, sample_id) in alignment.sample_ids.iter().enumerate() {
270 if !alignment.masks.iter().any(|mask| mask.present[idx]) {
271 return Err(DataError::Validation(format!(
272 "alignment sample `{sample_id}` is absent from every fused source"
273 )));
274 }
275 }
276 for block in blocks {
277 let mask = alignment
278 .masks
279 .iter()
280 .find(|mask| mask.source_id == block.source_id)
281 .expect("alignment sources were checked before presence validation");
282 let rows = row_maps
283 .get(&block.source_id)
284 .expect("source row map was created");
285 for (sample_id, present) in alignment.sample_ids.iter().zip(mask.present.iter()) {
286 let has_rows = rows.contains_key(sample_id);
287 if *present != has_rows {
288 return Err(DataError::Validation(format!(
289 "alignment presence for source `{}` sample `{sample_id}` is {present} but feature block rows are {}",
290 block.source_id,
291 if has_rows { "present" } else { "absent" }
292 )));
293 }
294 }
295 }
296 Ok(())
297}
298
299fn synthetic_observation_id(sample_id: &SampleId) -> Result<ObservationId> {
300 ObservationId::new(format!("fused.{}", sample_id.as_str()))
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306 use crate::alignment::{build_sample_alignment_plan, AlignmentMode, AlignmentPolicy};
307 use crate::ids::RepresentationId;
308 use serde_json::json;
309
310 fn block(
311 source_id: &str,
312 feature_names: &[&str],
313 rows: &[(&str, &str, Vec<serde_json::Value>)],
314 ) -> SourceFeatureBlock {
315 SourceFeatureBlock {
316 source_id: SourceId::new(source_id).unwrap(),
317 block: CoordinatorFeatureBlock {
318 feature_set_id: source_id.to_string(),
319 representation_id: RepresentationId::new("tabular_numeric").unwrap(),
320 feature_names: feature_names.iter().map(ToString::to_string).collect(),
321 observation_ids: rows
322 .iter()
323 .map(|(observation_id, _, _)| ObservationId::new(*observation_id).unwrap())
324 .collect(),
325 sample_ids: rows
326 .iter()
327 .map(|(_, sample_id, _)| SampleId::new(*sample_id).unwrap())
328 .collect(),
329 values: rows.iter().map(|(_, _, values)| values.clone()).collect(),
330 },
331 }
332 }
333
334 fn alignment(blocks: &[SourceFeatureBlock], mode: AlignmentMode) -> SampleAlignmentPlan {
335 let source_sets = blocks
336 .iter()
337 .map(source_sample_set_from_feature_block)
338 .collect::<Result<Vec<_>>>()
339 .unwrap();
340 build_sample_alignment_plan(&source_sets, &AlignmentPolicy { mode }).unwrap()
341 }
342
343 #[test]
344 fn fusion_broadcasts_singleton_source_to_reference_repetitions() {
345 let blocks = vec![
346 block(
347 "nir",
348 &["n0"],
349 &[
350 ("obs.S001.r1", "S001", vec![json!(1.0)]),
351 ("obs.S001.r2", "S001", vec![json!(2.0)]),
352 ("obs.S002.r1", "S002", vec![json!(3.0)]),
353 ],
354 ),
355 block(
356 "chem",
357 &["c0"],
358 &[
359 ("chem.S001", "S001", vec![json!(10.0)]),
360 ("chem.S002", "S002", vec![json!(20.0)]),
361 ],
362 ),
363 ];
364 let fused = fuse_feature_blocks(
365 "fused",
366 &blocks,
367 &alignment(&blocks, AlignmentMode::Inner),
368 &FeatureFusionPolicy::default(),
369 )
370 .unwrap();
371
372 assert_eq!(fused.feature_names, vec!["nir.n0", "chem.c0"]);
373 assert_eq!(fused.sample_ids, blocks[0].block.sample_ids);
374 assert_eq!(
375 fused.values,
376 vec![
377 vec![json!(1.0), json!(10.0)],
378 vec![json!(2.0), json!(10.0)],
379 vec![json!(3.0), json!(20.0)]
380 ]
381 );
382 }
383
384 #[test]
385 fn outer_fusion_creates_synthetic_rows_for_samples_missing_in_reference() {
386 let blocks = vec![
387 block("nir", &["n0"], &[("obs.S001.r1", "S001", vec![json!(1.0)])]),
388 block(
389 "chem",
390 &["c0"],
391 &[
392 ("chem.S001", "S001", vec![json!(10.0)]),
393 ("chem.S002", "S002", vec![json!(20.0)]),
394 ],
395 ),
396 ];
397 let fused = fuse_feature_blocks(
398 "fused",
399 &blocks,
400 &alignment(&blocks, AlignmentMode::Outer),
401 &FeatureFusionPolicy::default(),
402 )
403 .unwrap();
404
405 assert_eq!(
406 fused
407 .observation_ids
408 .iter()
409 .map(ToString::to_string)
410 .collect::<Vec<_>>(),
411 vec!["obs.S001.r1", "fused.S002"]
412 );
413 assert_eq!(
414 fused.values,
415 vec![
416 vec![json!(1.0), json!(10.0)],
417 vec![serde_json::Value::Null, json!(20.0)]
418 ]
419 );
420 }
421
422 #[test]
423 fn fusion_refuses_ambiguous_non_reference_repetitions() {
424 let blocks = vec![
425 block("chem", &["c0"], &[("chem.S001", "S001", vec![json!(10.0)])]),
426 block(
427 "nir",
428 &["n0"],
429 &[
430 ("obs.S001.r1", "S001", vec![json!(1.0)]),
431 ("obs.S001.r2", "S001", vec![json!(2.0)]),
432 ],
433 ),
434 ];
435 let err = fuse_feature_blocks(
436 "fused",
437 &blocks,
438 &alignment(&blocks, AlignmentMode::Inner),
439 &FeatureFusionPolicy::default(),
440 )
441 .unwrap_err();
442
443 assert!(err.to_string().contains("cannot broadcast"));
444 }
445
446 #[test]
447 fn fusion_refuses_duplicate_unnamespaced_feature_names() {
448 let blocks = vec![
449 block("nir", &["x"], &[("nir.S001", "S001", vec![json!(1.0)])]),
450 block("chem", &["x"], &[("chem.S001", "S001", vec![json!(10.0)])]),
451 ];
452 let err = fuse_feature_blocks(
453 "fused",
454 &blocks,
455 &alignment(&blocks, AlignmentMode::Inner),
456 &FeatureFusionPolicy {
457 namespace_columns: false,
458 },
459 )
460 .unwrap_err();
461
462 assert!(err.to_string().contains("duplicate feature"));
463 }
464
465 #[test]
466 fn fusion_refuses_alignment_presence_that_does_not_match_rows() {
467 let blocks = vec![
468 block("nir", &["n0"], &[("nir.S001", "S001", vec![json!(1.0)])]),
469 block("chem", &["c0"], &[("chem.S001", "S001", vec![json!(10.0)])]),
470 ];
471 let mut bad_alignment = alignment(&blocks, AlignmentMode::Inner);
472 bad_alignment.masks[1].present[0] = false;
473 let err = fuse_feature_blocks(
474 "fused",
475 &blocks,
476 &bad_alignment,
477 &FeatureFusionPolicy::default(),
478 )
479 .unwrap_err();
480
481 assert!(err.to_string().contains("alignment presence"));
482 }
483
484 #[test]
485 fn fusion_refuses_alignment_sample_absent_from_all_sources() {
486 let blocks = vec![
487 block("nir", &["n0"], &[("nir.S001", "S001", vec![json!(1.0)])]),
488 block("chem", &["c0"], &[("chem.S001", "S001", vec![json!(10.0)])]),
489 ];
490 let mut bad_alignment = alignment(&blocks, AlignmentMode::Inner);
491 bad_alignment.masks[0].present[0] = false;
492 bad_alignment.masks[1].present[0] = false;
493 let err = fuse_feature_blocks(
494 "fused",
495 &blocks,
496 &bad_alignment,
497 &FeatureFusionPolicy::default(),
498 )
499 .unwrap_err();
500
501 assert!(err.to_string().contains("absent from every fused source"));
502 }
503}