dag_ml_core/
relation.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::campaign::stable_json_fingerprint;
6use crate::error::{DagMlError, Result};
7use crate::fold::FoldSet;
8use crate::ids::{GroupId, ObservationId, SampleId, TargetId};
9use crate::policy::{LeakageUnitPolicy, SplitUnit};
10
11#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum FoldPartition {
14    Train,
15    Validation,
16}
17
18#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum EntityUnitLevel {
21    PhysicalSample,
22    SourceSample,
23    #[default]
24    Observation,
25    Combo,
26}
27
28#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
29#[serde(deny_unknown_fields)]
30pub struct SampleRelation {
31    #[serde(default)]
32    pub unit_level: EntityUnitLevel,
33    #[serde(default)]
34    pub unit_id: Option<String>,
35    pub observation_id: ObservationId,
36    pub sample_id: SampleId,
37    #[serde(default)]
38    pub source_id: Option<String>,
39    #[serde(default)]
40    pub rep_id: Option<String>,
41    #[serde(default)]
42    pub target_id: Option<TargetId>,
43    #[serde(default)]
44    pub group_id: Option<GroupId>,
45    #[serde(default)]
46    pub origin_sample_id: Option<SampleId>,
47    #[serde(default)]
48    pub derived_unit_id: Option<String>,
49    #[serde(default)]
50    pub component_observation_ids: Vec<ObservationId>,
51    #[serde(default)]
52    pub sample_influence_weight: Option<f64>,
53    #[serde(default)]
54    pub quality_flag: Option<String>,
55    #[serde(default)]
56    pub is_augmented: bool,
57    #[serde(default, skip_serializing_if = "is_false")]
58    pub excluded: bool,
59    // Metadata + tags carried so a `by_metadata` / `by_tag` branch view selector
60    // can match natively in the data provider. Skipped when empty so relation
61    // sets without them keep byte-identical fingerprints (existing fixtures and
62    // contracts stay unaffected). A non-empty value changes the replay
63    // fingerprint because it changes which samples a branch view selects.
64    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
65    pub metadata: BTreeMap<String, serde_json::Value>,
66    #[serde(default, skip_serializing_if = "Vec::is_empty")]
67    pub tags: Vec<String>,
68}
69
70fn is_false(value: &bool) -> bool {
71    !*value
72}
73
74impl SampleRelation {
75    pub fn new(observation_id: ObservationId, sample_id: SampleId) -> Self {
76        Self {
77            unit_level: EntityUnitLevel::Observation,
78            unit_id: None,
79            observation_id,
80            sample_id,
81            source_id: None,
82            rep_id: None,
83            target_id: None,
84            group_id: None,
85            origin_sample_id: None,
86            derived_unit_id: None,
87            component_observation_ids: Vec::new(),
88            sample_influence_weight: None,
89            quality_flag: None,
90            is_augmented: false,
91            excluded: false,
92            metadata: BTreeMap::new(),
93            tags: Vec::new(),
94        }
95    }
96
97    pub fn effective_unit_id(&self) -> Result<String> {
98        if let Some(unit_id) = non_empty_optional("unit_id", &self.observation_id, &self.unit_id)? {
99            return Ok(unit_id.to_string());
100        }
101
102        match self.unit_level {
103            EntityUnitLevel::PhysicalSample => Ok(self.sample_id.to_string()),
104            EntityUnitLevel::SourceSample => {
105                let source_id =
106                    non_empty_optional("source_id", &self.observation_id, &self.source_id)?
107                        .ok_or_else(|| {
108                            DagMlError::CampaignValidation(format!(
109                                "source-sample relation `{}` requires source_id",
110                                self.observation_id
111                            ))
112                        })?;
113                Ok(format!("{}::{source_id}", self.sample_id))
114            }
115            EntityUnitLevel::Observation => Ok(self.observation_id.to_string()),
116            EntityUnitLevel::Combo => {
117                let derived_unit_id = non_empty_optional(
118                    "derived_unit_id",
119                    &self.observation_id,
120                    &self.derived_unit_id,
121                )?
122                .ok_or_else(|| {
123                    DagMlError::CampaignValidation(format!(
124                        "combo relation `{}` requires derived_unit_id",
125                        self.observation_id
126                    ))
127                })?;
128                Ok(derived_unit_id.to_string())
129            }
130        }
131    }
132
133    fn validate(&self) -> Result<()> {
134        non_empty_optional("unit_id", &self.observation_id, &self.unit_id)?;
135        non_empty_optional("source_id", &self.observation_id, &self.source_id)?;
136        non_empty_optional(
137            "derived_unit_id",
138            &self.observation_id,
139            &self.derived_unit_id,
140        )?;
141        non_empty_optional("quality_flag", &self.observation_id, &self.quality_flag)?;
142        validate_optional_identifier("rep_id", &self.observation_id, &self.rep_id)?;
143
144        if let Some(weight) = self.sample_influence_weight {
145            if !weight.is_finite() || weight <= 0.0 {
146                return Err(DagMlError::CampaignValidation(format!(
147                    "relation `{}` has invalid sample_influence_weight",
148                    self.observation_id
149                )));
150            }
151        }
152
153        if self.unit_level != EntityUnitLevel::Combo && !self.component_observation_ids.is_empty() {
154            return Err(DagMlError::CampaignValidation(format!(
155                "relation `{}` has component_observation_ids but is not a combo",
156                self.observation_id
157            )));
158        }
159
160        self.effective_unit_id()?;
161        Ok(())
162    }
163}
164
165#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
166#[serde(deny_unknown_fields)]
167pub struct SampleRelationSet {
168    #[serde(default)]
169    pub records: Vec<SampleRelation>,
170}
171
172pub fn relation_set_fingerprint(relations: &SampleRelationSet) -> Result<String> {
173    relations.fingerprint()
174}
175
176#[derive(Clone, Debug, Serialize)]
177struct CanonicalRelationRecord {
178    effective_unit_id: String,
179    unit_level: EntityUnitLevel,
180    unit_id: Option<String>,
181    observation_id: ObservationId,
182    sample_id: SampleId,
183    source_id: Option<String>,
184    rep_id: Option<String>,
185    target_id: Option<TargetId>,
186    group_id: Option<GroupId>,
187    origin_sample_id: Option<SampleId>,
188    derived_unit_id: Option<String>,
189    component_observation_ids: Vec<ObservationId>,
190    sample_influence_weight: Option<f64>,
191    quality_flag: Option<String>,
192    is_augmented: bool,
193    // Skipped when false so `excluded=false` relation sets keep byte-identical
194    // fingerprints (and existing fixtures/contracts stay unaffected); an
195    // `excluded=true` row correctly changes the replay fingerprint because it
196    // changes the training set.
197    #[serde(default, skip_serializing_if = "is_false")]
198    excluded: bool,
199    // Skipped when empty so relations without metadata/tags keep byte-identical
200    // fingerprints; a non-empty value correctly changes the replay fingerprint
201    // because it changes which samples a `by_metadata` / `by_tag` branch view
202    // selects.
203    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
204    metadata: BTreeMap<String, serde_json::Value>,
205    #[serde(skip_serializing_if = "Vec::is_empty")]
206    tags: Vec<String>,
207}
208
209impl SampleRelationSet {
210    pub fn validate(&self) -> Result<()> {
211        let mut observations = BTreeSet::new();
212        let mut observation_samples = BTreeMap::<ObservationId, SampleId>::new();
213        let mut unit_ids = BTreeMap::<String, ObservationId>::new();
214        let mut sample_targets = BTreeMap::<SampleId, TargetId>::new();
215        let mut sample_groups = BTreeMap::<SampleId, GroupId>::new();
216        for record in &self.records {
217            record.validate()?;
218            if !observations.insert(&record.observation_id) {
219                return Err(DagMlError::CampaignValidation(format!(
220                    "duplicate observation relation `{}`",
221                    record.observation_id
222                )));
223            }
224            observation_samples.insert(record.observation_id.clone(), record.sample_id.clone());
225            let effective_unit_id = record.effective_unit_id()?;
226            if let Some(previous) =
227                unit_ids.insert(effective_unit_id.clone(), record.observation_id.clone())
228            {
229                return Err(DagMlError::CampaignValidation(format!(
230                    "relations `{previous}` and `{}` share effective unit id `{effective_unit_id}`",
231                    record.observation_id
232                )));
233            }
234            if let Some(target_id) = &record.target_id {
235                if let Some(previous) = sample_targets.get(&record.sample_id) {
236                    if previous != target_id {
237                        return Err(DagMlError::CampaignValidation(format!(
238                            "sample `{}` maps to multiple targets",
239                            record.sample_id
240                        )));
241                    }
242                } else {
243                    sample_targets.insert(record.sample_id.clone(), target_id.clone());
244                }
245            }
246            if let Some(group_id) = &record.group_id {
247                if let Some(previous) = sample_groups.get(&record.sample_id) {
248                    if previous != group_id {
249                        return Err(DagMlError::CampaignValidation(format!(
250                            "sample `{}` maps to multiple groups",
251                            record.sample_id
252                        )));
253                    }
254                } else {
255                    sample_groups.insert(record.sample_id.clone(), group_id.clone());
256                }
257            }
258        }
259        for record in &self.records {
260            validate_combo_record(record, &observation_samples)?;
261        }
262        Ok(())
263    }
264
265    pub fn fingerprint(&self) -> Result<String> {
266        self.validate()?;
267        let mut canonical = self
268            .records
269            .iter()
270            .map(|record| {
271                let effective_unit_id = record.effective_unit_id()?;
272                Ok(CanonicalRelationRecord {
273                    effective_unit_id,
274                    unit_level: record.unit_level,
275                    unit_id: record.unit_id.clone(),
276                    observation_id: record.observation_id.clone(),
277                    sample_id: record.sample_id.clone(),
278                    source_id: record.source_id.clone(),
279                    rep_id: record.rep_id.clone(),
280                    target_id: record.target_id.clone(),
281                    group_id: record.group_id.clone(),
282                    origin_sample_id: record.origin_sample_id.clone(),
283                    derived_unit_id: record.derived_unit_id.clone(),
284                    component_observation_ids: record.component_observation_ids.clone(),
285                    sample_influence_weight: record.sample_influence_weight,
286                    quality_flag: record.quality_flag.clone(),
287                    is_augmented: record.is_augmented,
288                    excluded: record.excluded,
289                    metadata: record.metadata.clone(),
290                    tags: record.tags.clone(),
291                })
292            })
293            .collect::<Result<Vec<_>>>()?;
294        canonical.sort_by(|left, right| {
295            (
296                left.effective_unit_id.as_str(),
297                left.observation_id.as_str(),
298                left.sample_id.as_str(),
299            )
300                .cmp(&(
301                    right.effective_unit_id.as_str(),
302                    right.observation_id.as_str(),
303                    right.sample_id.as_str(),
304                ))
305        });
306        stable_json_fingerprint(&canonical)
307    }
308
309    pub fn validate_against_fold_set(
310        &self,
311        fold_set: &FoldSet,
312        policy: &LeakageUnitPolicy,
313    ) -> Result<()> {
314        self.validate()?;
315        fold_set.validate()?;
316        policy.validate()?;
317
318        let universe = fold_set.sample_ids.iter().collect::<BTreeSet<_>>();
319        for record in &self.records {
320            if !universe.contains(&record.sample_id) {
321                return Err(DagMlError::CampaignValidation(format!(
322                    "relation `{}` references sample `{}` outside fold set",
323                    record.observation_id, record.sample_id
324                )));
325            }
326            if let Some(origin_sample_id) = &record.origin_sample_id {
327                if !universe.contains(origin_sample_id) {
328                    return Err(DagMlError::CampaignValidation(format!(
329                        "relation `{}` references origin sample `{}` outside fold set",
330                        record.observation_id, origin_sample_id
331                    )));
332                }
333            }
334            if policy.require_group_ids && record.group_id.is_none() {
335                return Err(DagMlError::CampaignValidation(format!(
336                    "relation `{}` is missing required group id",
337                    record.observation_id
338                )));
339            }
340        }
341
342        let sample_to_target = self.sample_targets();
343        let sample_to_group = self.sample_groups();
344        validate_fold_set_groups_match_relations(fold_set, &sample_to_group)?;
345
346        for fold in &fold_set.folds {
347            let partitions = fold
348                .train_sample_ids
349                .iter()
350                .map(|sample_id| (sample_id, FoldPartition::Train))
351                .chain(
352                    fold.validation_sample_ids
353                        .iter()
354                        .map(|sample_id| (sample_id, FoldPartition::Validation)),
355                )
356                .collect::<BTreeMap<_, _>>();
357
358            if policy.forbid_origin_cross_fold {
359                for record in &self.records {
360                    if let Some(origin_sample_id) = &record.origin_sample_id {
361                        let sample_partition =
362                            partitions.get(&record.sample_id).ok_or_else(|| {
363                                DagMlError::CampaignValidation(format!(
364                                    "fold `{}` does not contain sample `{}`",
365                                    fold.fold_id, record.sample_id
366                                ))
367                            })?;
368                        let origin_partition =
369                            partitions.get(origin_sample_id).ok_or_else(|| {
370                                DagMlError::CampaignValidation(format!(
371                                    "fold `{}` does not contain origin sample `{}`",
372                                    fold.fold_id, origin_sample_id
373                                ))
374                            })?;
375                        if sample_partition != origin_partition {
376                            return Err(DagMlError::CampaignValidation(format!(
377                                "fold `{}` leaks origin sample `{}` into {:?} sample `{}`",
378                                fold.fold_id, origin_sample_id, sample_partition, record.sample_id
379                            )));
380                        }
381                    }
382                }
383            }
384
385            match policy.split_unit {
386                SplitUnit::PhysicalSample | SplitUnit::Observation | SplitUnit::Sample => {}
387                SplitUnit::Target => validate_unit_partitions(
388                    &fold.fold_id.to_string(),
389                    "target",
390                    &partitions,
391                    &sample_to_target,
392                )?,
393                SplitUnit::Group => validate_unit_partitions(
394                    &fold.fold_id.to_string(),
395                    "group",
396                    &partitions,
397                    &sample_to_group,
398                )?,
399            }
400        }
401        Ok(())
402    }
403
404    pub fn sample_for_observation(&self, observation_id: &ObservationId) -> Option<&SampleId> {
405        self.records
406            .iter()
407            .find(|record| &record.observation_id == observation_id)
408            .map(|record| &record.sample_id)
409    }
410
411    pub fn target_for_sample(&self, sample_id: &SampleId) -> Option<&TargetId> {
412        self.records
413            .iter()
414            .find(|record| &record.sample_id == sample_id)
415            .and_then(|record| record.target_id.as_ref())
416    }
417
418    pub fn group_for_sample(&self, sample_id: &SampleId) -> Option<&GroupId> {
419        self.records
420            .iter()
421            .find(|record| &record.sample_id == sample_id)
422            .and_then(|record| record.group_id.as_ref())
423    }
424
425    pub fn observation_count_for_sample(&self, sample_id: &SampleId) -> usize {
426        self.records
427            .iter()
428            .filter(|record| &record.sample_id == sample_id)
429            .count()
430    }
431
432    /// Samples excluded from training. Exclusion is sample-local: a sample is
433    /// excluded if ANY of its relation rows carries `excluded == true`, so a
434    /// multi-source / repetition sample can never train through a sibling
435    /// non-excluded row. These samples are still validated and predicted.
436    pub fn excluded_sample_ids(&self) -> BTreeSet<SampleId> {
437        self.records
438            .iter()
439            .filter(|record| record.excluded)
440            .map(|record| record.sample_id.clone())
441            .collect()
442    }
443
444    pub fn sample_targets(&self) -> BTreeMap<SampleId, TargetId> {
445        self.records
446            .iter()
447            .filter_map(|record| {
448                record
449                    .target_id
450                    .as_ref()
451                    .map(|target_id| (record.sample_id.clone(), target_id.clone()))
452            })
453            .collect()
454    }
455
456    pub fn sample_groups(&self) -> BTreeMap<SampleId, GroupId> {
457        self.records
458            .iter()
459            .filter_map(|record| {
460                record
461                    .group_id
462                    .as_ref()
463                    .map(|group_id| (record.sample_id.clone(), group_id.clone()))
464            })
465            .collect()
466    }
467}
468
469fn non_empty_optional<'a>(
470    field: &str,
471    observation_id: &ObservationId,
472    value: &'a Option<String>,
473) -> Result<Option<&'a str>> {
474    if let Some(value) = value.as_deref() {
475        if value.trim().is_empty() {
476            return Err(DagMlError::CampaignValidation(format!(
477                "relation `{observation_id}` has empty {field}"
478            )));
479        }
480        Ok(Some(value))
481    } else {
482        Ok(None)
483    }
484}
485
486fn validate_optional_identifier(
487    field: &str,
488    observation_id: &ObservationId,
489    value: &Option<String>,
490) -> Result<()> {
491    let Some(value) = non_empty_optional(field, observation_id, value)? else {
492        return Ok(());
493    };
494    if value.len() > 128
495        || !value
496            .bytes()
497            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b':'))
498    {
499        return Err(DagMlError::CampaignValidation(format!(
500            "relation `{observation_id}` has invalid {field}"
501        )));
502    }
503    Ok(())
504}
505
506fn validate_combo_record(
507    record: &SampleRelation,
508    observation_samples: &BTreeMap<ObservationId, SampleId>,
509) -> Result<()> {
510    if record.unit_level != EntityUnitLevel::Combo {
511        return Ok(());
512    }
513    if record.component_observation_ids.is_empty() {
514        return Err(DagMlError::CampaignValidation(format!(
515            "combo relation `{}` has no component observations",
516            record.observation_id
517        )));
518    }
519    if record.derived_unit_id.is_none() {
520        return Err(DagMlError::CampaignValidation(format!(
521            "combo relation `{}` requires derived_unit_id",
522            record.observation_id
523        )));
524    }
525    if let Some(origin_sample_id) = &record.origin_sample_id {
526        if origin_sample_id != &record.sample_id {
527            return Err(DagMlError::CampaignValidation(format!(
528                "combo relation `{}` origin sample `{}` differs from sample `{}`",
529                record.observation_id, origin_sample_id, record.sample_id
530            )));
531        }
532    }
533
534    let mut components = BTreeSet::new();
535    for component_observation_id in &record.component_observation_ids {
536        if component_observation_id == &record.observation_id {
537            return Err(DagMlError::CampaignValidation(format!(
538                "combo relation `{}` cannot list itself as a component",
539                record.observation_id
540            )));
541        }
542        if !components.insert(component_observation_id) {
543            return Err(DagMlError::CampaignValidation(format!(
544                "combo relation `{}` repeats component observation `{}`",
545                record.observation_id, component_observation_id
546            )));
547        }
548        let component_sample = observation_samples
549            .get(component_observation_id)
550            .ok_or_else(|| {
551                DagMlError::CampaignValidation(format!(
552                    "combo relation `{}` references missing component observation `{}`",
553                    record.observation_id, component_observation_id
554                ))
555            })?;
556        if component_sample != &record.sample_id {
557            return Err(DagMlError::CampaignValidation(format!(
558                "combo relation `{}` component observation `{}` belongs to sample `{}` not `{}`",
559                record.observation_id, component_observation_id, component_sample, record.sample_id
560            )));
561        }
562    }
563    Ok(())
564}
565
566fn validate_fold_set_groups_match_relations(
567    fold_set: &FoldSet,
568    sample_to_group: &BTreeMap<SampleId, GroupId>,
569) -> Result<()> {
570    for (sample_id, fold_group) in &fold_set.sample_groups {
571        if let Some(relation_group) = sample_to_group.get(sample_id) {
572            if relation_group != fold_group {
573                return Err(DagMlError::CampaignValidation(format!(
574                    "sample `{sample_id}` has group `{relation_group}` in relations but `{fold_group}` in fold set"
575                )));
576            }
577        }
578    }
579    Ok(())
580}
581
582fn validate_unit_partitions<Unit: Ord + std::fmt::Display>(
583    fold_id: &str,
584    label: &str,
585    partitions: &BTreeMap<&SampleId, FoldPartition>,
586    sample_units: &BTreeMap<SampleId, Unit>,
587) -> Result<()> {
588    let mut unit_partitions = BTreeMap::<&Unit, FoldPartition>::new();
589    for (sample_id, partition) in partitions {
590        let Some(unit) = sample_units.get(*sample_id) else {
591            return Err(DagMlError::CampaignValidation(format!(
592                "fold `{fold_id}` sample `{sample_id}` is missing {label} id"
593            )));
594        };
595        if let Some(previous) = unit_partitions.insert(unit, *partition) {
596            if previous != *partition {
597                return Err(DagMlError::CampaignValidation(format!(
598                    "fold `{fold_id}` leaks {label} `{unit}` across train/validation"
599                )));
600            }
601        }
602    }
603    Ok(())
604}
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609    use crate::data::ExternalDataPlanEnvelope;
610    use crate::fold::{FoldAssignment, FoldPartitionMode};
611
612    fn sid(value: &str) -> SampleId {
613        SampleId::new(value).unwrap()
614    }
615
616    fn oid(value: &str) -> ObservationId {
617        ObservationId::new(value).unwrap()
618    }
619
620    fn tid(value: &str) -> TargetId {
621        TargetId::new(value).unwrap()
622    }
623
624    fn gid(value: &str) -> GroupId {
625        GroupId::new(value).unwrap()
626    }
627
628    fn fold_set() -> FoldSet {
629        FoldSet {
630            id: "outer".to_string(),
631            sample_ids: vec![sid("s1"), sid("s2"), sid("s3"), sid("s4")],
632            folds: vec![
633                FoldAssignment {
634                    fold_id: crate::ids::FoldId::new("fold:0").unwrap(),
635                    train_sample_ids: vec![sid("s3"), sid("s4")],
636                    validation_sample_ids: vec![sid("s1"), sid("s2")],
637                    metadata: BTreeMap::new(),
638                },
639                FoldAssignment {
640                    fold_id: crate::ids::FoldId::new("fold:1").unwrap(),
641                    train_sample_ids: vec![sid("s1"), sid("s2")],
642                    validation_sample_ids: vec![sid("s3"), sid("s4")],
643                    metadata: BTreeMap::new(),
644                },
645            ],
646            sample_groups: BTreeMap::new(),
647            partition_mode: FoldPartitionMode::Partition,
648        }
649    }
650
651    fn relation(observation: &str, sample: &str, target: &str, group: &str) -> SampleRelation {
652        let mut relation = SampleRelation::new(oid(observation), sid(sample));
653        relation.target_id = Some(tid(target));
654        relation.group_id = Some(gid(group));
655        relation
656    }
657
658    fn source_relation(observation: &str, sample: &str, source: &str, rep: &str) -> SampleRelation {
659        let mut relation = relation(observation, sample, "target:sample", "group:sample");
660        relation.source_id = Some(source.to_string());
661        relation.rep_id = Some(rep.to_string());
662        relation
663    }
664
665    #[test]
666    fn repeated_observations_validate_at_sample_split_unit() {
667        let relations = SampleRelationSet {
668            records: vec![
669                relation("obs:1a", "s1", "t1", "g1"),
670                relation("obs:1b", "s1", "t1", "g1"),
671                relation("obs:2a", "s2", "t2", "g2"),
672                relation("obs:3a", "s3", "t3", "g3"),
673                relation("obs:4a", "s4", "t4", "g4"),
674            ],
675        };
676
677        relations
678            .validate_against_fold_set(&fold_set(), &LeakageUnitPolicy::default())
679            .unwrap();
680    }
681
682    #[test]
683    fn repeated_observations_validate_at_physical_sample_split_unit() {
684        let relations = SampleRelationSet {
685            records: vec![
686                relation("obs:1a", "s1", "t1", "g1"),
687                relation("obs:1b", "s1", "t1", "g1"),
688                relation("obs:2a", "s2", "t2", "g2"),
689                relation("obs:3a", "s3", "t3", "g3"),
690                relation("obs:4a", "s4", "t4", "g4"),
691            ],
692        };
693        let policy = LeakageUnitPolicy {
694            split_unit: SplitUnit::PhysicalSample,
695            ..LeakageUnitPolicy::default()
696        };
697
698        relations
699            .validate_against_fold_set(&fold_set(), &policy)
700            .unwrap();
701    }
702
703    #[test]
704    fn asymmetric_multisource_repetitions_and_combo_validate_as_relations() {
705        let mut combo = relation(
706            "obs:s1.combo.a0.b0.c0",
707            "s1",
708            "target:sample",
709            "group:sample",
710        );
711        combo.unit_level = EntityUnitLevel::Combo;
712        combo.source_id = Some("combo".to_string());
713        combo.derived_unit_id = Some("combo:s1:a0:b0:c0".to_string());
714        combo.origin_sample_id = Some(sid("s1"));
715        combo.component_observation_ids =
716            vec![oid("obs:s1.A.0"), oid("obs:s1.B.0"), oid("obs:s1.C.0")];
717        combo.sample_influence_weight = Some(1.0);
718        combo.quality_flag = Some("ok".to_string());
719
720        let relations = SampleRelationSet {
721            records: vec![
722                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
723                source_relation("obs:s1.A.1", "s1", "A", "rep:1"),
724                source_relation("obs:s1.B.0", "s1", "B", "rep:0"),
725                source_relation("obs:s1.B.1", "s1", "B", "rep:1"),
726                source_relation("obs:s1.B.2", "s1", "B", "rep:2"),
727                source_relation("obs:s1.C.0", "s1", "C", "rep:0"),
728                source_relation("obs:s1.C.1", "s1", "C", "rep:1"),
729                combo,
730            ],
731        };
732
733        relations.validate().unwrap();
734        assert_eq!(
735            relations.sample_for_observation(&oid("obs:s1.combo.a0.b0.c0")),
736            Some(&sid("s1"))
737        );
738    }
739
740    #[test]
741    fn combo_components_cannot_cross_sample_boundary() {
742        let mut combo = relation("obs:s1.combo", "s1", "target:sample", "group:sample");
743        combo.unit_level = EntityUnitLevel::Combo;
744        combo.derived_unit_id = Some("combo:s1".to_string());
745        combo.component_observation_ids = vec![oid("obs:s1.A.0"), oid("obs:s2.B.0")];
746
747        let relations = SampleRelationSet {
748            records: vec![
749                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
750                source_relation("obs:s2.B.0", "s2", "B", "rep:0"),
751                combo,
752            ],
753        };
754
755        assert!(relations.validate().is_err());
756    }
757
758    #[test]
759    fn relation_fingerprint_is_order_stable_and_provenance_sensitive() {
760        let left = SampleRelationSet {
761            records: vec![
762                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
763                source_relation("obs:s1.B.0", "s1", "B", "rep:0"),
764            ],
765        };
766        let right = SampleRelationSet {
767            records: vec![
768                source_relation("obs:s1.B.0", "s1", "B", "rep:0"),
769                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
770            ],
771        };
772        assert_eq!(left.fingerprint().unwrap(), right.fingerprint().unwrap());
773
774        let mut changed = left.clone();
775        changed.records[0].rep_id = Some("rep:1".to_string());
776        assert_ne!(left.fingerprint().unwrap(), changed.fingerprint().unwrap());
777    }
778
779    #[test]
780    fn excluded_bit_changes_fingerprint_but_only_when_true() {
781        let base = SampleRelationSet {
782            records: vec![
783                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
784                source_relation("obs:s2.A.0", "s2", "A", "rep:0"),
785            ],
786        };
787
788        // excluded=false is byte-identical to the default (skip_serializing_if):
789        // existing fixtures and contracts stay unaffected.
790        let mut explicit_false = base.clone();
791        explicit_false.records[0].excluded = false;
792        assert_eq!(
793            base.fingerprint().unwrap(),
794            explicit_false.fingerprint().unwrap()
795        );
796
797        // excluded=true changes the training set, so it MUST change the
798        // replay fingerprint (else a different exclusion mask could false-hit a
799        // cached bundle).
800        let mut excluded = base.clone();
801        excluded.records[0].excluded = true;
802        assert_ne!(base.fingerprint().unwrap(), excluded.fingerprint().unwrap());
803    }
804
805    #[test]
806    fn metadata_and_tags_change_fingerprint_but_only_when_non_empty() {
807        let base = SampleRelationSet {
808            records: vec![
809                source_relation("obs:s1.A.0", "s1", "A", "rep:0"),
810                source_relation("obs:s2.A.0", "s2", "A", "rep:0"),
811            ],
812        };
813
814        // Empty metadata/tags are skip-serialized, so an explicit empty value is
815        // byte-identical to the default: existing fixtures/contracts unaffected.
816        let mut explicit_empty = base.clone();
817        explicit_empty.records[0].metadata = BTreeMap::new();
818        explicit_empty.records[0].tags = Vec::new();
819        assert_eq!(
820            base.fingerprint().unwrap(),
821            explicit_empty.fingerprint().unwrap()
822        );
823
824        // Metadata changes which samples a `by_metadata` view selects, so it
825        // MUST change the replay fingerprint.
826        let mut with_metadata = base.clone();
827        with_metadata.records[0]
828            .metadata
829            .insert("group".to_string(), serde_json::json!("A"));
830        assert_ne!(
831            base.fingerprint().unwrap(),
832            with_metadata.fingerprint().unwrap()
833        );
834
835        // Same for tags (used by `by_tag` views).
836        let mut with_tags = base.clone();
837        with_tags.records[0].tags = vec!["clean".to_string()];
838        assert_ne!(
839            base.fingerprint().unwrap(),
840            with_tags.fingerprint().unwrap()
841        );
842    }
843
844    #[test]
845    fn old_relation_json_defaults_to_observation_unit() {
846        let relation: SampleRelation = serde_json::from_value(serde_json::json!({
847            "observation_id": "obs:legacy",
848            "sample_id": "s1",
849            "target_id": "t1",
850            "group_id": "g1",
851            "source_id": "legacy",
852            "is_augmented": false
853        }))
854        .unwrap();
855
856        assert_eq!(relation.unit_level, EntityUnitLevel::Observation);
857        assert!(relation.rep_id.is_none());
858        assert!(relation.component_observation_ids.is_empty());
859        SampleRelationSet {
860            records: vec![relation],
861        }
862        .validate()
863        .unwrap();
864    }
865
866    #[test]
867    fn relation_contracts_reject_unknown_fields_at_every_nesting_level() {
868        let relation = serde_json::json!({
869            "observation_id": "obs:strict",
870            "sample_id": "sample:strict"
871        });
872
873        let mut unknown_set_field = serde_json::json!({"records": [relation.clone()]});
874        unknown_set_field.as_object_mut().unwrap().insert(
875            "unexpected_contract_field".to_string(),
876            serde_json::json!(true),
877        );
878        assert!(serde_json::from_value::<SampleRelationSet>(unknown_set_field).is_err());
879
880        let mut unknown_record_field = relation.clone();
881        unknown_record_field.as_object_mut().unwrap().insert(
882            "unexpected_contract_field".to_string(),
883            serde_json::json!(true),
884        );
885        assert!(serde_json::from_value::<SampleRelation>(unknown_record_field.clone()).is_err());
886
887        let envelope = serde_json::json!({
888            "schema_version": 1,
889            "schema_fingerprint": "0".repeat(64),
890            "plan_fingerprint": "1".repeat(64),
891            "coordinator_relations": {"records": [unknown_record_field]}
892        });
893        assert!(serde_json::from_value::<ExternalDataPlanEnvelope>(envelope).is_err());
894
895        let envelope = serde_json::json!({
896            "schema_version": 1,
897            "schema_fingerprint": "0".repeat(64),
898            "plan_fingerprint": "1".repeat(64),
899            "coordinator_relations": {
900                "records": [relation.clone()],
901                "unexpected_contract_field": true
902            }
903        });
904        assert!(serde_json::from_value::<ExternalDataPlanEnvelope>(envelope).is_err());
905
906        let accepted: SampleRelationSet = serde_json::from_value(serde_json::json!({
907            "records": [{
908                "observation_id": "obs:strict",
909                "sample_id": "sample:strict",
910                "metadata": {"unexpected_contract_field": "opaque metadata remains open"},
911                "tags": ["strict"]
912            }]
913        }))
914        .unwrap();
915        assert_eq!(
916            accepted.records[0].metadata["unexpected_contract_field"],
917            "opaque metadata remains open"
918        );
919    }
920
921    #[test]
922    fn relation_validation_rejects_invalid_new_fields() {
923        let mut invalid_rep = source_relation("obs:s1.A.0", "s1", "A", "rep/0");
924        assert!(invalid_rep.validate().is_err());
925
926        invalid_rep.rep_id = Some("rep:0".to_string());
927        invalid_rep.sample_influence_weight = Some(0.0);
928        assert!(invalid_rep.validate().is_err());
929
930        invalid_rep.sample_influence_weight = Some(1.0);
931        invalid_rep.quality_flag = Some(" ".to_string());
932        assert!(invalid_rep.validate().is_err());
933    }
934
935    #[test]
936    fn target_split_refuses_shared_target_across_fold_boundary() {
937        let relations = SampleRelationSet {
938            records: vec![
939                relation("obs:1", "s1", "same_target", "g1"),
940                relation("obs:2", "s2", "t2", "g2"),
941                relation("obs:3", "s3", "same_target", "g3"),
942                relation("obs:4", "s4", "t4", "g4"),
943            ],
944        };
945        let policy = LeakageUnitPolicy {
946            split_unit: SplitUnit::Target,
947            ..LeakageUnitPolicy::default()
948        };
949
950        assert!(relations
951            .validate_against_fold_set(&fold_set(), &policy)
952            .is_err());
953    }
954
955    #[test]
956    fn augmentation_origin_cannot_cross_train_validation_boundary() {
957        let mut generated = relation("obs:aug", "s3", "t3", "g3");
958        generated.origin_sample_id = Some(sid("s1"));
959        generated.is_augmented = true;
960        let relations = SampleRelationSet {
961            records: vec![
962                relation("obs:1", "s1", "t1", "g1"),
963                relation("obs:2", "s2", "t2", "g2"),
964                generated,
965                relation("obs:4", "s4", "t4", "g4"),
966            ],
967        };
968
969        assert!(relations
970            .validate_against_fold_set(&fold_set(), &LeakageUnitPolicy::default())
971            .is_err());
972    }
973}
dag_ml_core/relation.rs

dag_ml_core/
relation.rs
