Skip to main content

dag_ml_data_core/
fingerprint.rs

1use sha2::{Digest, Sha256};
2
3use crate::error::Result;
4use crate::model::DatasetSchema;
5use crate::plan::DataPlan;
6use crate::relation::{FoldSet, SampleRelationTable};
7
8pub fn schema_fingerprint(schema: &DatasetSchema) -> Result<String> {
9    let mut canonical = schema.clone();
10    canonical.validate()?;
11    canonical.sample_ids.sort();
12    canonical
13        .sources
14        .sort_by(|left, right| left.id.cmp(&right.id));
15    canonical
16        .groups
17        .sort_by(|left, right| left.id.cmp(&right.id));
18    canonical
19        .folds
20        .sort_by(|left, right| left.id.cmp(&right.id));
21
22    let json = serde_json::to_vec(&canonical)?;
23    let digest = Sha256::digest(json);
24    Ok(to_hex(&digest))
25}
26
27pub fn data_plan_fingerprint(plan: &DataPlan) -> Result<String> {
28    plan.validate()?;
29    let json = serde_json::to_vec(plan)?;
30    let digest = Sha256::digest(json);
31    Ok(to_hex(&digest))
32}
33
34pub fn sample_relation_fingerprint(relations: &SampleRelationTable) -> Result<String> {
35    let mut canonical = relations.clone();
36    canonical.validate()?;
37    canonical.rows.sort_by(|left, right| {
38        left.observation_id
39            .cmp(&right.observation_id)
40            .then_with(|| left.sample_id.cmp(&right.sample_id))
41            .then_with(|| left.source_id.cmp(&right.source_id))
42    });
43
44    let json = serde_json::to_vec(&canonical)?;
45    let digest = Sha256::digest(json);
46    Ok(to_hex(&digest))
47}
48
49pub fn fold_set_fingerprint(fold_set: &FoldSet) -> Result<String> {
50    let mut canonical = fold_set.clone();
51    canonical.validate()?;
52    canonical.sample_ids.sort();
53    canonical
54        .folds
55        .sort_by(|left, right| left.fold_id.cmp(&right.fold_id));
56    for fold in &mut canonical.folds {
57        fold.train_sample_ids.sort();
58        fold.validation_sample_ids.sort();
59    }
60
61    let mut value = serde_json::to_value(&canonical)?;
62    remove_empty_fold_set_maps(&mut value);
63    let json = serde_json::to_vec(&value)?;
64    let digest = Sha256::digest(json);
65    Ok(to_hex(&digest))
66}
67
68fn remove_empty_fold_set_maps(value: &mut serde_json::Value) {
69    let Some(object) = value.as_object_mut() else {
70        return;
71    };
72    if object
73        .get("sample_groups")
74        .and_then(serde_json::Value::as_object)
75        .is_some_and(serde_json::Map::is_empty)
76    {
77        object.remove("sample_groups");
78    }
79    let Some(folds) = object
80        .get_mut("folds")
81        .and_then(serde_json::Value::as_array_mut)
82    else {
83        return;
84    };
85    for fold in folds {
86        let Some(fold_object) = fold.as_object_mut() else {
87            continue;
88        };
89        if fold_object
90            .get("metadata")
91            .and_then(serde_json::Value::as_object)
92            .is_some_and(serde_json::Map::is_empty)
93        {
94            fold_object.remove("metadata");
95        }
96    }
97}
98
99fn to_hex(bytes: &[u8]) -> String {
100    let mut out = String::with_capacity(bytes.len() * 2);
101    for byte in bytes {
102        use std::fmt::Write;
103        write!(&mut out, "{byte:02x}").expect("writing to string cannot fail");
104    }
105    out
106}
107
108#[cfg(test)]
109mod tests {
110    use std::collections::BTreeMap;
111
112    use crate::ids::{GroupId, RepresentationId, SampleId, SourceId, TypeId};
113    use crate::model::{
114        AxisKind, AxisSpec, CoordinateDType, CoordinateSpec, CoordinateValues, DatasetSchema,
115        FoldSpec, GroupKind, GroupSpec, RepresentationSpec, SourceDescriptor, SourceGranularity,
116    };
117    use crate::plan::DataPlan;
118    use crate::relation::{FoldAssignment, FoldSet};
119
120    use super::{
121        data_plan_fingerprint, fold_set_fingerprint, sample_relation_fingerprint,
122        schema_fingerprint,
123    };
124
125    const SHARED_FOLD_SET_FINGERPRINT: &str =
126        "54d3185d6c628ef0df848828a8d8ae650222a283a78bbd3ab3bc2256f222c05c";
127
128    fn representation(id: &str) -> RepresentationSpec {
129        RepresentationSpec {
130            id: RepresentationId::new(id).unwrap(),
131            type_id: TypeId::new("table").unwrap(),
132            rank: Some(2),
133            axes: vec![
134                AxisSpec {
135                    name: "sample".to_string(),
136                    kind: AxisKind::Sample,
137                    unit: None,
138                    size: Some(2),
139                    variable: false,
140                    coordinate: None,
141                },
142                AxisSpec {
143                    name: "feature".to_string(),
144                    kind: AxisKind::Feature,
145                    unit: None,
146                    size: Some(1),
147                    variable: false,
148                    coordinate: None,
149                },
150            ],
151            container: "dataframe".to_string(),
152            dtype: Some("float32".to_string()),
153            sparse: false,
154            ragged: false,
155            signal_type: None,
156        }
157    }
158
159    fn source(id: &str) -> SourceDescriptor {
160        SourceDescriptor {
161            id: SourceId::new(id).unwrap(),
162            name: id.to_string(),
163            type_id: TypeId::new("table").unwrap(),
164            modality: "metadata".to_string(),
165            native_representation: representation("tabular"),
166            sample_key: "sample_id".to_string(),
167            granularity: SourceGranularity::PerSample,
168            schema: BTreeMap::new(),
169            tags: BTreeMap::new(),
170            shape_contract: None,
171        }
172    }
173
174    #[test]
175    fn fingerprint_is_independent_of_source_order() {
176        let mut left = DatasetSchema {
177            dataset_id: "d".to_string(),
178            sample_ids: vec![SampleId::new("s2").unwrap(), SampleId::new("s1").unwrap()],
179            sources: vec![source("b"), source("a")],
180            targets: BTreeMap::new(),
181            metadata: BTreeMap::new(),
182            metadata_schema: None,
183            groups: vec![
184                GroupSpec {
185                    id: GroupId::new("g.b").unwrap(),
186                    kind: GroupKind::Batch,
187                    column: "batch_b".to_string(),
188                    source_id: Some(SourceId::new("b").unwrap()),
189                    strict: false,
190                    metadata: BTreeMap::new(),
191                },
192                GroupSpec {
193                    id: GroupId::new("g.a").unwrap(),
194                    kind: GroupKind::RepetitionGroup,
195                    column: "sample_id".to_string(),
196                    source_id: Some(SourceId::new("a").unwrap()),
197                    strict: true,
198                    metadata: BTreeMap::new(),
199                },
200            ],
201            folds: vec![
202                FoldSpec {
203                    id: "fold.b".to_string(),
204                    group_id: Some(GroupId::new("g.b").unwrap()),
205                    split_column: Some("fold_b".to_string()),
206                    metadata: BTreeMap::new(),
207                },
208                FoldSpec {
209                    id: "fold.a".to_string(),
210                    group_id: Some(GroupId::new("g.a").unwrap()),
211                    split_column: Some("fold_a".to_string()),
212                    metadata: BTreeMap::new(),
213                },
214            ],
215        };
216        let mut right = left.clone();
217        right.sources.reverse();
218        right.sample_ids.reverse();
219        right.groups.reverse();
220        right.folds.reverse();
221
222        assert_eq!(
223            schema_fingerprint(&left).unwrap(),
224            schema_fingerprint(&right).unwrap()
225        );
226
227        left.dataset_id = "different".to_string();
228        assert_ne!(
229            schema_fingerprint(&left).unwrap(),
230            schema_fingerprint(&right).unwrap()
231        );
232    }
233
234    #[test]
235    fn data_plan_fingerprint_is_stable() {
236        let plan: DataPlan = serde_json::from_str(include_str!(
237            "../../../examples/fixtures/oof_campaign/expected_data_plan_nir_to_tabular.json"
238        ))
239        .unwrap();
240
241        assert_eq!(
242            data_plan_fingerprint(&plan).unwrap(),
243            data_plan_fingerprint(&plan).unwrap()
244        );
245    }
246
247    #[test]
248    fn sample_relation_fingerprint_is_stable() {
249        let relations: crate::relation::SampleRelationTable = serde_json::from_str(include_str!(
250            "../../../examples/fixtures/oof_campaign/sample_relations_grouped_augmented.json"
251        ))
252        .unwrap();
253
254        assert_eq!(
255            sample_relation_fingerprint(&relations).unwrap(),
256            sample_relation_fingerprint(&relations).unwrap()
257        );
258    }
259
260    #[test]
261    fn fold_set_fingerprint_is_independent_of_ordering() {
262        let mut left = FoldSet {
263            id: "cv.partition".to_string(),
264            sample_ids: vec![
265                SampleId::new("s3").unwrap(),
266                SampleId::new("s2").unwrap(),
267                SampleId::new("s1").unwrap(),
268            ],
269            folds: vec![
270                FoldAssignment {
271                    fold_id: "fold1".to_string(),
272                    train_sample_ids: vec![
273                        SampleId::new("s2").unwrap(),
274                        SampleId::new("s1").unwrap(),
275                    ],
276                    validation_sample_ids: vec![SampleId::new("s3").unwrap()],
277                    metadata: BTreeMap::new(),
278                },
279                FoldAssignment {
280                    fold_id: "fold0".to_string(),
281                    train_sample_ids: vec![SampleId::new("s3").unwrap()],
282                    validation_sample_ids: vec![
283                        SampleId::new("s2").unwrap(),
284                        SampleId::new("s1").unwrap(),
285                    ],
286                    metadata: BTreeMap::new(),
287                },
288            ],
289            sample_groups: BTreeMap::new(),
290        };
291        let mut right = left.clone();
292        right.sample_ids.reverse();
293        right.folds.reverse();
294        for fold in &mut right.folds {
295            fold.train_sample_ids.reverse();
296            fold.validation_sample_ids.reverse();
297        }
298
299        assert_eq!(
300            fold_set_fingerprint(&left).unwrap(),
301            fold_set_fingerprint(&right).unwrap()
302        );
303
304        left.id = "cv.partition.changed".to_string();
305        assert_ne!(
306            fold_set_fingerprint(&left).unwrap(),
307            fold_set_fingerprint(&right).unwrap()
308        );
309    }
310
311    fn coordinate_schema(coordinate: Option<CoordinateSpec>) -> DatasetSchema {
312        let mut repr = representation("tabular");
313        // The feature axis (size 1) carries the coordinate.
314        repr.axes[1].coordinate = coordinate;
315        let mut descriptor = source("a");
316        descriptor.native_representation = repr;
317        DatasetSchema {
318            dataset_id: "coords".to_string(),
319            sample_ids: vec![SampleId::new("s1").unwrap()],
320            sources: vec![descriptor],
321            targets: BTreeMap::new(),
322            metadata: BTreeMap::new(),
323            metadata_schema: None,
324            groups: Vec::new(),
325            folds: Vec::new(),
326        }
327    }
328
329    #[test]
330    fn schema_fingerprint_reflects_axis_coordinates() {
331        let explicit = CoordinateSpec {
332            dtype: CoordinateDType::Categorical,
333            ordered: false,
334            values: CoordinateValues::Explicit {
335                values: vec![serde_json::Value::from("R")],
336            },
337        };
338        let grid = CoordinateSpec {
339            dtype: CoordinateDType::Numeric,
340            ordered: true,
341            values: CoordinateValues::RegularGrid {
342                start: 400.0,
343                step: 2.0,
344            },
345        };
346
347        let bare = schema_fingerprint(&coordinate_schema(None)).unwrap();
348        let with_explicit = schema_fingerprint(&coordinate_schema(Some(explicit.clone()))).unwrap();
349        let with_grid = schema_fingerprint(&coordinate_schema(Some(grid))).unwrap();
350
351        // Stable for a fixed coordinate spec.
352        assert_eq!(
353            with_explicit,
354            schema_fingerprint(&coordinate_schema(Some(explicit))).unwrap()
355        );
356        // Coordinates participate in the fingerprint and the two shapes differ.
357        assert_ne!(bare, with_explicit);
358        assert_ne!(bare, with_grid);
359        assert_ne!(with_explicit, with_grid);
360    }
361
362    #[test]
363    fn shared_fold_set_fixture_fingerprint_is_locked() {
364        let fixture = include_str!("../../../examples/fixtures/shared/fold_set_cv_partition.json");
365        let fold_set = serde_json::from_str::<FoldSet>(fixture).unwrap();
366
367        assert_eq!(
368            fold_set_fingerprint(&fold_set).unwrap(),
369            SHARED_FOLD_SET_FINGERPRINT
370        );
371    }
372}