1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{DataError, Result};
6use crate::ids::{GroupId, ObservationId, OriginId, RepetitionId, SampleId, SourceId, TargetId};
7
8#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
9pub struct AugmentationMetadata {
10 pub transform_id: String,
11 #[serde(default, skip_serializing_if = "Option::is_none")]
12 pub params_fingerprint: Option<String>,
13 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
14 pub metadata: BTreeMap<String, serde_json::Value>,
15}
16
17impl AugmentationMetadata {
18 pub fn validate(&self) -> Result<()> {
19 if self.transform_id.trim().is_empty() {
20 return Err(DataError::Validation(
21 "augmentation metadata transform_id is empty".to_string(),
22 ));
23 }
24 if let Some(params_fingerprint) = &self.params_fingerprint {
25 if params_fingerprint.len() != 64
26 || !params_fingerprint
27 .bytes()
28 .all(|byte| byte.is_ascii_hexdigit())
29 {
30 return Err(DataError::Validation(format!(
31 "augmentation `{}` params_fingerprint must be a 64-character hex digest",
32 self.transform_id
33 )));
34 }
35 }
36 for key in self.metadata.keys() {
37 if key.trim().is_empty() {
38 return Err(DataError::Validation(format!(
39 "augmentation `{}` metadata contains an empty key",
40 self.transform_id
41 )));
42 }
43 }
44 Ok(())
45 }
46}
47
48#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
49pub struct SampleRelation {
50 pub observation_id: ObservationId,
51 pub sample_id: SampleId,
52 pub source_id: Option<SourceId>,
53 pub target_id: Option<TargetId>,
54 pub group_id: Option<GroupId>,
55 pub origin_id: Option<OriginId>,
56 pub repetition_id: Option<RepetitionId>,
57 #[serde(default)]
58 pub augmented: bool,
59 #[serde(default)]
60 pub excluded: bool,
61 #[serde(default)]
62 pub metadata: BTreeMap<String, serde_json::Value>,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
64 pub augmentation: Option<AugmentationMetadata>,
65}
66
67#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
68pub struct SampleRelationTable {
69 pub rows: Vec<SampleRelation>,
70}
71
72#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
73pub struct FoldAssignment {
74 pub fold_id: String,
75 pub train_sample_ids: Vec<SampleId>,
76 pub validation_sample_ids: Vec<SampleId>,
77 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
78 pub metadata: BTreeMap<String, serde_json::Value>,
79}
80
81#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
87pub struct FoldSet {
88 pub id: String,
89 pub sample_ids: Vec<SampleId>,
90 pub folds: Vec<FoldAssignment>,
91 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
92 pub sample_groups: BTreeMap<SampleId, GroupId>,
93}
94
95impl FoldSet {
96 pub fn validate(&self) -> Result<()> {
97 if self.id.trim().is_empty() {
98 return Err(DataError::Validation("fold set id is empty".to_string()));
99 }
100 if self.sample_ids.is_empty() {
101 return Err(DataError::Validation(
102 "fold set contains no samples".to_string(),
103 ));
104 }
105 if self.folds.is_empty() {
106 return Err(DataError::Validation(
107 "fold set contains no folds".to_string(),
108 ));
109 }
110 let universe = unique_samples("fold set sample_ids", &self.sample_ids)?;
111 if !self.sample_groups.is_empty() {
112 for sample_id in self.sample_groups.keys() {
113 if !universe.contains(sample_id) {
114 return Err(DataError::Validation(format!(
115 "sample group map references unknown sample `{sample_id}`"
116 )));
117 }
118 }
119 }
120
121 let mut fold_ids = BTreeSet::new();
122 let mut validation_counts = self
123 .sample_ids
124 .iter()
125 .cloned()
126 .map(|sample_id| (sample_id, 0usize))
127 .collect::<BTreeMap<_, _>>();
128
129 for fold in &self.folds {
130 if fold.fold_id.trim().is_empty() {
131 return Err(DataError::Validation(
132 "fold assignment id is empty".to_string(),
133 ));
134 }
135 if !fold_ids.insert(&fold.fold_id) {
136 return Err(DataError::Validation(format!(
137 "duplicate fold id `{}`",
138 fold.fold_id
139 )));
140 }
141 let train = unique_samples(
142 &format!("fold `{}` train_sample_ids", fold.fold_id),
143 &fold.train_sample_ids,
144 )?;
145 let validation = unique_samples(
146 &format!("fold `{}` validation_sample_ids", fold.fold_id),
147 &fold.validation_sample_ids,
148 )?;
149 if validation.is_empty() {
150 return Err(DataError::Validation(format!(
151 "fold `{}` has no validation samples",
152 fold.fold_id
153 )));
154 }
155 for sample_id in train.union(&validation) {
156 if !universe.contains(sample_id) {
157 return Err(DataError::Validation(format!(
158 "fold `{}` references unknown sample `{}`",
159 fold.fold_id, sample_id
160 )));
161 }
162 }
163 if let Some(sample_id) = train.intersection(&validation).next() {
164 return Err(DataError::Validation(format!(
165 "fold `{}` has train/validation overlap at sample `{}`",
166 fold.fold_id, sample_id
167 )));
168 }
169 for sample_id in &validation {
170 *validation_counts
171 .get_mut(*sample_id)
172 .expect("validation sample is in universe") += 1;
173 }
174 validate_group_boundary(&fold.fold_id, &train, &validation, &self.sample_groups)?;
175 }
176
177 for (sample_id, count) in validation_counts {
178 if count != 1 {
179 return Err(DataError::Validation(format!(
180 "sample `{sample_id}` appears in validation {count} time(s), expected exactly once"
181 )));
182 }
183 }
184 Ok(())
185 }
186}
187
188impl SampleRelationTable {
189 pub fn validate(&self) -> Result<()> {
190 if self.rows.is_empty() {
191 return Err(DataError::Validation(
192 "sample relation table contains no rows".to_string(),
193 ));
194 }
195
196 let mut observation_ids = BTreeSet::new();
197 let mut origin_ids = BTreeSet::new();
198 let mut sample_groups = BTreeMap::<&SampleId, &GroupId>::new();
199 for row in &self.rows {
200 if !observation_ids.insert(&row.observation_id) {
201 return Err(DataError::Validation(format!(
202 "duplicate observation id `{}`",
203 row.observation_id
204 )));
205 }
206 if let Some(group_id) = &row.group_id {
207 if let Some(previous_group_id) = sample_groups.insert(&row.sample_id, group_id) {
208 if previous_group_id != group_id {
209 return Err(DataError::Validation(format!(
210 "sample `{}` appears with conflicting groups `{}` and `{}`",
211 row.sample_id, previous_group_id, group_id
212 )));
213 }
214 }
215 }
216 if row.augmented && row.origin_id.is_none() {
217 return Err(DataError::Validation(format!(
218 "augmented observation `{}` has no origin_id",
219 row.observation_id
220 )));
221 }
222 if !row.augmented && row.origin_id.is_some() {
223 return Err(DataError::Validation(format!(
224 "non-augmented observation `{}` declares origin_id",
225 row.observation_id
226 )));
227 }
228 if let Some(augmentation) = &row.augmentation {
229 augmentation.validate()?;
230 if !row.augmented {
231 return Err(DataError::Validation(format!(
232 "non-augmented observation `{}` declares augmentation metadata",
233 row.observation_id
234 )));
235 }
236 }
237 if let Some(origin_id) = &row.origin_id {
238 if origin_id.as_str() == row.observation_id.as_str() {
239 return Err(DataError::Validation(format!(
240 "observation `{}` cannot be its own origin",
241 row.observation_id
242 )));
243 }
244 origin_ids.insert(origin_id);
245 }
246 }
247
248 for origin_id in origin_ids {
249 if !observation_ids
250 .iter()
251 .any(|observation_id| observation_id.as_str() == origin_id.as_str())
252 {
253 return Err(DataError::Validation(format!(
254 "origin `{origin_id}` is not present as an observation"
255 )));
256 }
257 }
258
259 Ok(())
260 }
261
262 pub fn sample_groups(&self) -> BTreeMap<SampleId, GroupId> {
263 self.rows
264 .iter()
265 .filter_map(|row| {
266 row.group_id
267 .as_ref()
268 .map(|group_id| (row.sample_id.clone(), group_id.clone()))
269 })
270 .collect()
271 }
272
273 pub fn validate_fold_set(&self, fold_set: &FoldSet) -> Result<()> {
274 self.validate()?;
275 fold_set.validate()?;
276
277 let relation_sample_ids = self
278 .rows
279 .iter()
280 .map(|row| row.sample_id.clone())
281 .collect::<BTreeSet<_>>();
282 let fold_sample_ids = fold_set.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
283 if let Some(sample_id) = relation_sample_ids.difference(&fold_sample_ids).next() {
284 return Err(DataError::Validation(format!(
285 "fold set `{}` is missing relation sample `{sample_id}`",
286 fold_set.id
287 )));
288 }
289 if let Some(sample_id) = fold_sample_ids.difference(&relation_sample_ids).next() {
290 return Err(DataError::Validation(format!(
291 "fold set `{}` references sample `{sample_id}` absent from sample relations",
292 fold_set.id
293 )));
294 }
295
296 let relation_groups = self.sample_groups();
297 for (sample_id, fold_group_id) in &fold_set.sample_groups {
298 if let Some(relation_group_id) = relation_groups.get(sample_id) {
299 if relation_group_id != fold_group_id {
300 return Err(DataError::Validation(format!(
301 "fold set `{}` group `{}` for sample `{sample_id}` conflicts with relation group `{}`",
302 fold_set.id, fold_group_id, relation_group_id
303 )));
304 }
305 }
306 }
307
308 for fold in &fold_set.folds {
309 let train = fold.train_sample_ids.iter().collect::<BTreeSet<_>>();
310 let validation = fold.validation_sample_ids.iter().collect::<BTreeSet<_>>();
311 validate_group_boundary(&fold.fold_id, &train, &validation, &relation_groups)?;
312 validate_origin_boundary(&fold.fold_id, &train, &validation, self)?;
313 }
314 Ok(())
315 }
316}
317
318fn unique_samples<'a>(label: &str, samples: &'a [SampleId]) -> Result<BTreeSet<&'a SampleId>> {
319 let mut seen = BTreeSet::new();
320 for sample_id in samples {
321 if !seen.insert(sample_id) {
322 return Err(DataError::Validation(format!(
323 "{label} contains duplicate sample `{sample_id}`"
324 )));
325 }
326 }
327 Ok(seen)
328}
329
330fn validate_group_boundary(
331 fold_id: &str,
332 train: &BTreeSet<&SampleId>,
333 validation: &BTreeSet<&SampleId>,
334 sample_groups: &BTreeMap<SampleId, GroupId>,
335) -> Result<()> {
336 if sample_groups.is_empty() {
337 return Ok(());
338 }
339 let train_groups = train
340 .iter()
341 .filter_map(|sample_id| sample_groups.get(*sample_id))
342 .collect::<BTreeSet<_>>();
343 for sample_id in validation {
344 if let Some(group_id) = sample_groups.get(*sample_id) {
345 if train_groups.contains(group_id) {
346 return Err(DataError::RelationBoundaryViolation {
347 kind: "group",
348 detail: format!(
349 "fold `{fold_id}` leaks group `{group_id}` across train/validation"
350 ),
351 });
352 }
353 }
354 }
355 Ok(())
356}
357
358fn validate_origin_boundary(
359 fold_id: &str,
360 train: &BTreeSet<&SampleId>,
361 validation: &BTreeSet<&SampleId>,
362 relations: &SampleRelationTable,
363) -> Result<()> {
364 let origin_sample_ids = relations
365 .rows
366 .iter()
367 .map(|row| (row.observation_id.as_str(), &row.sample_id))
368 .collect::<BTreeMap<_, _>>();
369 for row in &relations.rows {
370 let Some(origin_id) = &row.origin_id else {
371 continue;
372 };
373 let Some(origin_sample_id) = origin_sample_ids.get(origin_id.as_str()) else {
374 continue;
375 };
376 if train.contains(&row.sample_id) && validation.contains(origin_sample_id) {
377 return Err(DataError::RelationBoundaryViolation {
378 kind: "origin",
379 detail: format!(
380 "fold `{fold_id}` leaks augmentation origin `{origin_id}` across train/validation"
381 ),
382 });
383 }
384 if validation.contains(&row.sample_id) && train.contains(origin_sample_id) {
385 return Err(DataError::RelationBoundaryViolation {
386 kind: "origin",
387 detail: format!(
388 "fold `{fold_id}` leaks augmentation origin `{origin_id}` across train/validation"
389 ),
390 });
391 }
392 }
393 Ok(())
394}
395
396#[cfg(test)]
397mod tests {
398 use super::*;
399
400 fn obs(value: &str) -> ObservationId {
401 ObservationId::new(value).unwrap()
402 }
403
404 fn sample(value: &str) -> SampleId {
405 SampleId::new(value).unwrap()
406 }
407
408 fn origin(value: &str) -> OriginId {
409 OriginId::new(value).unwrap()
410 }
411
412 fn row(observation_id: &str, sample_id: &str) -> SampleRelation {
413 SampleRelation {
414 observation_id: obs(observation_id),
415 sample_id: sample(sample_id),
416 source_id: None,
417 target_id: None,
418 group_id: None,
419 origin_id: None,
420 repetition_id: None,
421 augmented: false,
422 excluded: false,
423 metadata: BTreeMap::new(),
424 augmentation: None,
425 }
426 }
427
428 fn fold_set() -> FoldSet {
429 FoldSet {
430 id: "cv.group.safe".to_string(),
431 sample_ids: vec![sample("s1"), sample("s2")],
432 folds: vec![
433 FoldAssignment {
434 fold_id: "fold:0".to_string(),
435 train_sample_ids: vec![sample("s2")],
436 validation_sample_ids: vec![sample("s1")],
437 metadata: BTreeMap::new(),
438 },
439 FoldAssignment {
440 fold_id: "fold:1".to_string(),
441 train_sample_ids: vec![sample("s1")],
442 validation_sample_ids: vec![sample("s2")],
443 metadata: BTreeMap::new(),
444 },
445 ],
446 sample_groups: BTreeMap::new(),
447 }
448 }
449
450 #[test]
451 fn validates_group_and_origin_relations() {
452 let mut base = row("obs1", "s1");
453 base.group_id = Some(GroupId::new("g1").unwrap());
454 let mut augmented = row("obs1_aug", "s1");
455 augmented.group_id = Some(GroupId::new("g1").unwrap());
456 augmented.origin_id = Some(origin("obs1"));
457 augmented.augmented = true;
458 augmented.augmentation = Some(AugmentationMetadata {
459 transform_id: "snv.augment".to_string(),
460 params_fingerprint: Some("a".repeat(64)),
461 metadata: BTreeMap::new(),
462 });
463
464 let table = SampleRelationTable {
465 rows: vec![base, augmented],
466 };
467
468 table.validate().unwrap();
469 assert_eq!(
470 table.sample_groups().get(&sample("s1")),
471 Some(&GroupId::new("g1").unwrap())
472 );
473 }
474
475 #[test]
476 fn rejects_duplicate_observations() {
477 let table = SampleRelationTable {
478 rows: vec![row("obs1", "s1"), row("obs1", "s2")],
479 };
480
481 assert!(table.validate().is_err());
482 }
483
484 #[test]
485 fn rejects_augmented_rows_without_known_origin() {
486 let mut augmented = row("obs1_aug", "s1");
487 augmented.origin_id = Some(origin("missing"));
488 augmented.augmented = true;
489 let table = SampleRelationTable {
490 rows: vec![augmented],
491 };
492
493 assert!(table.validate().is_err());
494 }
495
496 #[test]
497 fn rejects_invalid_augmentation_metadata() {
498 let mut augmented = row("obs1_aug", "s1");
499 augmented.origin_id = Some(origin("obs1"));
500 augmented.augmented = true;
501 augmented.augmentation = Some(AugmentationMetadata {
502 transform_id: "snv.augment".to_string(),
503 params_fingerprint: Some("not-a-fingerprint".to_string()),
504 metadata: BTreeMap::new(),
505 });
506 let table = SampleRelationTable {
507 rows: vec![row("obs1", "s1"), augmented],
508 };
509
510 assert!(table.validate().is_err());
511 }
512
513 #[test]
514 fn validates_fold_set_against_relation_groups() {
515 let mut s1 = row("obs1", "s1");
516 s1.group_id = Some(GroupId::new("g1").unwrap());
517 let mut s2 = row("obs2", "s2");
518 s2.group_id = Some(GroupId::new("g2").unwrap());
519 let relations = SampleRelationTable { rows: vec![s1, s2] };
520
521 relations.validate_fold_set(&fold_set()).unwrap();
522 }
523
524 #[test]
525 fn rejects_fold_set_with_missing_validation_assignment() {
526 let mut fold_set = fold_set();
527 fold_set.folds = vec![FoldAssignment {
528 fold_id: "fold:0".to_string(),
529 train_sample_ids: vec![sample("s1")],
530 validation_sample_ids: vec![sample("s2")],
531 metadata: BTreeMap::new(),
532 }];
533
534 let error = fold_set.validate().unwrap_err();
535 assert!(error.to_string().contains("expected exactly once"));
536 }
537
538 #[test]
539 fn rejects_fold_set_with_repeated_validation_assignment() {
540 let mut fold_set = fold_set();
541 fold_set.folds = vec![
542 FoldAssignment {
543 fold_id: "fold:0".to_string(),
544 train_sample_ids: vec![sample("s2")],
545 validation_sample_ids: vec![sample("s1")],
546 metadata: BTreeMap::new(),
547 },
548 FoldAssignment {
549 fold_id: "fold:1".to_string(),
550 train_sample_ids: vec![sample("s2")],
551 validation_sample_ids: vec![sample("s1")],
552 metadata: BTreeMap::new(),
553 },
554 ];
555
556 let error = fold_set.validate().unwrap_err();
557 assert!(error.to_string().contains("expected exactly once"));
558 }
559
560 #[test]
561 fn rejects_fold_set_that_splits_relation_group() {
562 let mut s1 = row("obs1", "s1");
563 s1.group_id = Some(GroupId::new("g1").unwrap());
564 let mut s2 = row("obs2", "s2");
565 s2.group_id = Some(GroupId::new("g1").unwrap());
566 let relations = SampleRelationTable { rows: vec![s1, s2] };
567
568 let error = relations.validate_fold_set(&fold_set()).unwrap_err();
569 assert_eq!(error.category(), "data");
570 assert_eq!(error.code(), "relation_boundary_violation");
571 assert_eq!(error.error_code(), 0x0002_0002);
572 assert!(error.to_string().contains("leaks group"));
573 }
574
575 #[test]
576 fn rejects_fold_set_that_splits_augmentation_origin() {
577 let origin_row = row("obs1", "s1");
578 let mut augmented = row("obs1_aug", "s2");
579 augmented.augmented = true;
580 augmented.origin_id = Some(origin("obs1"));
581 let relations = SampleRelationTable {
582 rows: vec![origin_row, augmented],
583 };
584
585 let error = relations.validate_fold_set(&fold_set()).unwrap_err();
586 assert!(error.to_string().contains("augmentation origin"));
587 }
588
589 #[test]
590 fn rejects_fold_set_with_relation_sample_mismatch() {
591 let relations = SampleRelationTable {
592 rows: vec![row("obs1", "s1")],
593 };
594
595 let error = relations.validate_fold_set(&fold_set()).unwrap_err();
596 assert!(error.to_string().contains("absent from sample relations"));
597 }
598
599 #[test]
600 fn grouped_augmented_fixture_validates() {
601 let table: SampleRelationTable = serde_json::from_str(include_str!(
602 "../../../examples/fixtures/oof_campaign/sample_relations_grouped_augmented.json"
603 ))
604 .unwrap();
605
606 table.validate().unwrap();
607 assert_eq!(table.sample_groups().len(), 2);
608 }
609}