1use std::collections::HashMap;
7use std::path::Path;
8
9use crate::pipeline::PipelineConfig;
10
11#[derive(serde::Deserialize)]
12pub struct DomainSchema {
13 pub entities: Vec<EntitySchema>,
14 pub hn_types: Vec<HnSchema>,
15}
16
17fn default_entity_weight() -> f64 {
18 1.0
19}
20
21fn infer_identifier_col(entity: &EntitySchema) -> Option<String> {
24 let entity_id_name = format!("{}_id", entity.name);
25 let col_names: Vec<&str> = entity
26 .columns
27 .iter()
28 .filter_map(|c| c.get("name").and_then(|v| v.as_str()))
29 .collect();
30 if col_names.contains(&entity_id_name.as_str()) {
31 Some(entity_id_name)
32 } else if col_names.contains(&"id") {
33 Some("id".to_string())
34 } else {
35 col_names
36 .iter()
37 .find(|n| n.ends_with("_id"))
38 .map(|n| (*n).to_string())
39 }
40}
41
42#[derive(serde::Deserialize)]
43pub struct EntitySchema {
44 pub name: String,
45 pub columns: Vec<serde_json::Value>,
46 #[serde(default)]
47 pub fk_remaps: Vec<serde_json::Value>,
48 #[serde(default = "default_entity_weight")]
59 pub weight: f64,
60}
61
62#[derive(serde::Deserialize)]
63pub struct HnSchema {
64 pub entity_type: String,
65 pub config_json: String,
66}
67
68#[derive(Debug, Clone)]
69struct DifficultySettings {
70 singleton: f64,
71 doublet: f64,
72 noise_types: &'static [&'static str],
84 passes: usize,
108}
109
110const DIFFICULTY_MAP: [(&str, DifficultySettings); 3] = [
111 (
112 "light",
113 DifficultySettings {
114 singleton: 0.50,
115 doublet: 0.30,
116 noise_types: &["names", "dates"],
117 passes: 1,
118 },
119 ),
120 (
121 "medium",
122 DifficultySettings {
123 singleton: 0.30,
124 doublet: 0.40,
125 noise_types: &["typo", "names", "dates", "identifiers"],
126 passes: 3,
138 },
139 ),
140 (
141 "hell",
142 DifficultySettings {
143 singleton: 0.10,
144 doublet: 0.20,
145 noise_types: &[
151 "typo_aggressive",
152 "visual",
153 "unicode_pollution",
154 "names",
155 "dates",
156 "identifiers",
157 "addresses",
158 "companies",
159 "extra",
160 ],
161 passes: 8,
173 },
174 ),
175];
176
177pub fn default_singleton_master_fraction(difficulty: &str) -> f64 {
178 difficulty_settings(difficulty).singleton
179}
180
181fn difficulty_settings(difficulty: &str) -> DifficultySettings {
182 DIFFICULTY_MAP
183 .iter()
184 .find(|(name, _)| *name == difficulty)
185 .map(|(_, s)| s.clone())
186 .unwrap_or_else(|| DIFFICULTY_MAP[1].1.clone())
187}
188
189pub fn chrono_now() -> String {
191 let start = std::time::SystemTime::now()
192 .duration_since(std::time::UNIX_EPOCH)
193 .unwrap_or_default();
194 format!("{:x}", start.as_secs())
195}
196
197#[allow(clippy::too_many_arguments)]
208pub fn deterministic_run_id(
209 domain: &str,
210 size: usize,
211 seed: u64,
212 difficulty: &str,
213 hard_neg_ratio: f64,
214 singleton_master_fraction: f64,
215 locale: &str,
216 only_entity: Option<&str>,
217 chunk_size: Option<usize>,
218) -> String {
219 use std::hash::{Hash, Hasher};
220 let mut hasher = std::collections::hash_map::DefaultHasher::new();
221 domain.hash(&mut hasher);
222 size.hash(&mut hasher);
223 seed.hash(&mut hasher);
224 difficulty.hash(&mut hasher);
225 hard_neg_ratio.to_bits().hash(&mut hasher);
226 singleton_master_fraction.to_bits().hash(&mut hasher);
227 locale.hash(&mut hasher);
228 only_entity.unwrap_or("").hash(&mut hasher);
232 chunk_size.unwrap_or(0).hash(&mut hasher);
240 format!("{}_{:x}", domain, hasher.finish())
241}
242
243pub fn load_schema(domain: &str, schemas_dir: &Path) -> Result<DomainSchema, String> {
248 let path = schemas_dir.join(format!("{domain}.json"));
249 let available = available_domain_names(schemas_dir);
255 if !available.iter().any(|d| d == domain) {
256 let hint = if available.is_empty() {
257 "no schemas found".to_string()
258 } else {
259 available.join(", ")
260 };
261 return Err(format!(
262 "schema file not found for domain '{domain}' at {path:?}. \
263 Available domains ({hint})"
264 ));
265 }
266 let data =
267 std::fs::read_to_string(&path).map_err(|e| format!("cannot read schema {path:?}: {e}"))?;
268 serde_json::from_str(&data).map_err(|e| format!("cannot parse schema {domain}.json: {e}"))
269}
270
271fn available_domain_names(dir: &Path) -> Vec<String> {
273 match std::fs::read_dir(dir) {
274 Ok(entries) => entries
275 .filter_map(|e| e.ok())
276 .filter(|e| {
277 e.path()
278 .extension()
279 .map(|ext| ext == "json")
280 .unwrap_or(false)
281 })
282 .filter_map(|e| {
283 e.path()
284 .file_stem()
285 .map(|s| s.to_string_lossy().to_string())
286 })
287 .collect(),
288 Err(_) => Vec::new(),
289 }
290}
291
292#[allow(clippy::too_many_arguments)]
297pub fn build_pipeline_config(
298 domain: &str,
299 size: usize,
300 seed: u64,
301 difficulty: &str,
302 hard_neg_ratio: f64,
303 singleton_master_fraction: f64,
304 schema: &DomainSchema,
305 run_id: &str,
306 output_format: &str,
307 graph_enabled: bool,
308 graph_format: &str,
309 only_entity: Option<&str>,
310) -> Result<PipelineConfig, String> {
311 if size < 10 {
312 return Err(format!("size must be >= 10, got {size}"));
313 }
314 if schema.entities.is_empty() {
315 return Err(format!("schema for domain '{domain}' has no entities"));
316 }
317 if !(0.0..=1.0).contains(&singleton_master_fraction) {
318 return Err(format!(
319 "singleton_master_fraction must be in [0.0, 1.0], got {singleton_master_fraction}"
320 ));
321 }
322 if let Some(target) = only_entity
323 && !schema.entities.iter().any(|e| e.name == target)
324 {
325 let available: Vec<&str> = schema.entities.iter().map(|e| e.name.as_str()).collect();
326 return Err(format!(
327 "only_entity '{target}' not found in domain '{domain}'; available entities: {}",
328 available.join(", ")
329 ));
330 }
331 let ds = difficulty_settings(difficulty);
332 let total = size;
333
334 let n_singleton = (total as f64 * singleton_master_fraction) as usize;
335 let n_doublet_float = total as f64 * ds.doublet;
336 let mut n_doublet = n_doublet_float as usize;
337 if !n_doublet.is_multiple_of(2) {
338 n_doublet -= 1;
339 }
340 if n_singleton + n_doublet > total {
341 return Err(format!(
342 "singleton_master_fraction {singleton_master_fraction} leaves no room for this \
343 difficulty's doublet share ({:.2}); reduce singleton_master_fraction",
344 ds.doublet
345 ));
346 }
347 let mut n_triplet = total - n_singleton - n_doublet;
348 let r = n_triplet % 3;
349 if r != 0 {
350 n_triplet -= r;
351 }
352 let total_unique = n_singleton + n_doublet / 2 + n_triplet / 3;
353 let n_duplicates = total.max(total_unique) - total_unique;
354
355 let active_entities: Vec<&EntitySchema> = match only_entity {
362 Some(target) => schema
363 .entities
364 .iter()
365 .filter(|e| e.name == target)
366 .collect(),
367 None => schema.entities.iter().collect(),
368 };
369
370 let total_ratio: f64 = active_entities.iter().map(|e| e.weight).sum::<f64>();
371 let raw_floats: Vec<(&str, f64)> = active_entities
372 .iter()
373 .map(|e| {
374 (
375 e.name.as_str(),
376 total_unique as f64 * e.weight / total_ratio,
377 )
378 })
379 .collect();
380
381 let mut floor_map: HashMap<&str, usize> = HashMap::new();
382 for (name, r) in &raw_floats {
383 floor_map.insert(name, r.max(2.0) as usize);
384 }
385 let floor_sum: usize = floor_map.values().sum();
386 let need = total_unique.max(floor_sum) - floor_sum;
387 if need > 0 {
388 let mut remainders: Vec<(&str, f64)> = raw_floats
389 .iter()
390 .map(|(n, r)| (*n, r - r.floor()))
391 .collect();
392 remainders.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
393 for (name, _) in remainders.iter().take(need) {
394 *floor_map.get_mut(name).unwrap_or(&mut 0) += 1;
395 }
396 }
397
398 let dup_ratios: Vec<(&str, f64)> = active_entities
399 .iter()
400 .map(|e| {
401 (
402 e.name.as_str(),
403 *floor_map.get(e.name.as_str()).unwrap_or(&2) as f64 / total_unique as f64,
404 )
405 })
406 .collect();
407 let mut dup_floor: HashMap<&str, usize> = HashMap::new();
408 for (name, r) in &dup_ratios {
409 dup_floor.insert(name, (n_duplicates as f64 * r) as usize);
410 }
411 let dup_sum: usize = dup_floor.values().sum();
412 let dup_need = n_duplicates.max(dup_sum) - dup_sum;
413 if dup_need > 0 {
414 let mut remainders: Vec<(&str, f64)> = dup_ratios
415 .iter()
416 .map(|(n, r)| (*n, r - r.floor()))
417 .collect();
418 remainders.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
419 for (name, _) in remainders.iter().take(dup_need) {
420 *dup_floor.get_mut(name).unwrap_or(&mut 0) += 1;
421 }
422 }
423
424 let noise_count = ds.noise_types.len();
425 let noise_weights: Vec<f64> = vec![1.0 / noise_count as f64; noise_count];
426
427 let mut entity_plans = Vec::new();
428 for entity in &active_entities {
429 let n_base = *floor_map.get(entity.name.as_str()).unwrap_or(&2);
430 let n_dup = *dup_floor.get(entity.name.as_str()).unwrap_or(&0);
431
432 let mut noise_entries = Vec::new();
433 if n_dup > 0 {
434 let mut counts: Vec<usize> = noise_weights
435 .iter()
436 .map(|w| (w * n_dup as f64) as usize)
437 .collect();
438 let count_sum: usize = counts.iter().sum();
439 if count_sum < n_dup {
440 *counts.last_mut().unwrap_or(&mut 0) += n_dup - count_sum;
441 }
442 for (i, noise_type) in ds.noise_types.iter().enumerate() {
443 if counts[i] == 0 {
444 continue;
445 }
446 noise_entries.push(serde_json::json!({
447 "noise_type": noise_type,
448 "columns": [],
449 "count": counts[i],
450 }));
451 }
452 }
453
454 let identifier_col = infer_identifier_col(entity);
455
456 let columns_json = serde_json::to_string(&entity.columns)
457 .map_err(|e| format!("serialize columns: {e}"))?;
458
459 entity_plans.push(serde_json::json!({
460 "name": entity.name,
461 "n_base": n_base,
462 "n_dup": n_dup,
463 "identifier_col": identifier_col,
464 "columns_json": columns_json,
465 "noise_types": noise_entries,
466 "fk_remaps": entity.fk_remaps,
467 }));
468 }
469
470 if let Some(target) = only_entity {
477 let target_schema = schema
478 .entities
479 .iter()
480 .find(|e| e.name == target)
481 .expect("only_entity validated to exist above");
482 let mut pool_target_names: Vec<&str> = target_schema
483 .fk_remaps
484 .iter()
485 .filter_map(|r| r.get("target_entity").and_then(|v| v.as_str()))
486 .filter(|&t| t != target)
487 .collect();
488 pool_target_names.sort_unstable();
489 pool_target_names.dedup();
490
491 let pool_n = crate::pipeline::FK_POOL_CAP.min(total_unique.max(2));
492 for name in pool_target_names {
493 let Some(entity) = schema.entities.iter().find(|e| e.name == name) else {
494 continue;
495 };
496 let identifier_col = infer_identifier_col(entity);
497 let columns_json = serde_json::to_string(&entity.columns)
498 .map_err(|e| format!("serialize columns: {e}"))?;
499 entity_plans.push(serde_json::json!({
500 "name": entity.name,
501 "n_base": pool_n,
502 "n_dup": 0,
503 "identifier_col": identifier_col,
504 "columns_json": columns_json,
505 "noise_types": [],
506 "fk_remaps": entity.fk_remaps,
507 "pool_only": true,
508 }));
509 }
510 }
511
512 let active_hn_types: Vec<&HnSchema> = schema
516 .hn_types
517 .iter()
518 .filter(|hn| only_entity.is_none_or(|t| hn.entity_type == t))
519 .collect();
520
521 let n_hard_neg = (n_duplicates as f64 * hard_neg_ratio * 0.125) as usize;
534 let hn_per_type = n_hard_neg / active_hn_types.len().max(1);
535 let hard_neg_types: Vec<serde_json::Value> = active_hn_types
536 .iter()
537 .map(|hn| {
538 serde_json::json!({
539 "entity_type": hn.entity_type,
540 "config_json": hn.config_json,
541 "count": hn_per_type,
542 })
543 })
544 .collect();
545
546 let config = serde_json::json!({
547 "domain": domain,
548 "size": size,
549 "seed": seed,
550 "difficulty": difficulty,
551 "output_format": output_format,
552 "run_id": run_id,
553 "entity_plans": entity_plans,
554 "hard_neg_types": hard_neg_types,
555 "hard_neg_ratio": hard_neg_ratio,
556 "noise_passes": ds.passes,
557 "graph_enabled": graph_enabled,
558 "graph_format": graph_format,
559 });
560
561 serde_json::from_value(config).map_err(|e| format!("build PipelineConfig: {e}"))
562}
563
564#[cfg(test)]
565mod tests {
566 use super::*;
567
568 fn schemas_dir() -> std::path::PathBuf {
569 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas")
570 }
571
572 fn kyc_schema() -> DomainSchema {
573 load_schema("kyc", &schemas_dir()).expect("load kyc.json")
574 }
575
576 fn aviation_schema() -> DomainSchema {
577 load_schema("aviation", &schemas_dir()).expect("load aviation.json")
578 }
579
580 #[test]
581 fn test_load_schema_known_domain() {
582 let schema = kyc_schema();
583 assert_eq!(schema.entities.len(), 2);
584 assert!(schema.entities.iter().any(|e| e.name == "natural_person"));
585 assert!(!schema.hn_types.is_empty());
586 }
587
588 #[test]
589 fn test_load_schema_unknown_domain_lists_available() {
590 let err = load_schema("not-a-real-domain", &schemas_dir())
591 .err()
592 .unwrap();
593 assert!(err.contains("not-a-real-domain"));
594 assert!(err.contains("kyc"));
595 }
596
597 #[test]
598 fn test_load_schema_case_sensitive() {
599 let err = load_schema("KYC", &schemas_dir()).err().unwrap();
602 assert!(err.contains("KYC"));
603 }
604
605 #[test]
606 fn test_build_pipeline_config_basic() {
607 let schema = kyc_schema();
608 let config = build_pipeline_config(
609 "kyc", 1000, 42, "medium", 0.1, 0.3, &schema, "kyc_test", "parquet", false, "parquet",
610 None,
611 )
612 .expect("build config");
613 assert_eq!(config.domain, "kyc");
614 assert_eq!(config.size, 1000);
615 assert_eq!(config.entity_plans.len(), 2);
616 assert!(config.entity_plans.iter().all(|p| p.n_base >= 2));
618 assert!(!config.hard_neg_types.is_empty());
619 }
620
621 #[test]
622 fn test_build_pipeline_config_rejects_size_below_10() {
623 let schema = kyc_schema();
624 let err = build_pipeline_config(
625 "kyc", 5, 42, "medium", 0.1, 0.3, &schema, "kyc_test", "parquet", false, "parquet",
626 None,
627 )
628 .unwrap_err();
629 assert!(err.contains("size must be >= 10"));
630 }
631
632 #[test]
633 fn test_build_pipeline_config_rejects_invalid_singleton_fraction() {
634 let schema = kyc_schema();
635 let err = build_pipeline_config(
636 "kyc", 1000, 42, "medium", 0.1, 1.5, &schema, "kyc_test", "parquet", false, "parquet",
637 None,
638 )
639 .unwrap_err();
640 assert!(err.contains("singleton_master_fraction"));
641 }
642
643 #[test]
644 fn test_build_pipeline_config_deterministic() {
645 let schema = kyc_schema();
646 let build = || {
647 build_pipeline_config(
648 "kyc", 1000, 42, "hell", 0.1, 0.2, &schema, "kyc_test", "parquet", false,
649 "parquet", None,
650 )
651 .expect("build config")
652 };
653 let a = build();
654 let b = build();
655 assert_eq!(a.entity_plans.len(), b.entity_plans.len());
656 for (pa, pb) in a.entity_plans.iter().zip(b.entity_plans.iter()) {
657 assert_eq!(pa.n_base, pb.n_base);
658 assert_eq!(pa.noise_types.len(), pb.noise_types.len());
659 }
660 }
661
662 #[test]
666 fn test_build_pipeline_config_only_entity_no_fk() {
667 let schema = kyc_schema();
668 let config = build_pipeline_config(
669 "kyc",
670 1000,
671 42,
672 "medium",
673 0.1,
674 0.3,
675 &schema,
676 "kyc_test",
677 "parquet",
678 false,
679 "parquet",
680 Some("natural_person"),
681 )
682 .expect("build config");
683 assert_eq!(config.entity_plans.len(), 1);
684 let plan = &config.entity_plans[0];
685 assert_eq!(plan.name, "natural_person");
686 assert!(!plan.pool_only);
687 assert!(plan.n_base > 100);
691 assert!(
692 config
693 .hard_neg_types
694 .iter()
695 .all(|hn| hn.entity_type == "natural_person")
696 );
697 }
698
699 #[test]
706 fn test_build_pipeline_config_only_entity_with_fk_pool() {
707 let schema = aviation_schema();
708 let config = build_pipeline_config(
709 "aviation",
710 10_000,
711 42,
712 "medium",
713 0.1,
714 0.3,
715 &schema,
716 "aviation_test",
717 "parquet",
718 false,
719 "parquet",
720 Some("aircraft"),
721 )
722 .expect("build config");
723 assert_eq!(config.entity_plans.len(), 2);
725 let target = config
726 .entity_plans
727 .iter()
728 .find(|p| p.name == "aircraft")
729 .expect("aircraft plan present");
730 assert!(!target.pool_only);
731 assert!(target.n_base > 1_000);
732 let pool = config
733 .entity_plans
734 .iter()
735 .find(|p| p.name == "airline")
736 .expect("airline pool-only plan present");
737 assert!(pool.pool_only);
738 assert!(pool.noise_types.is_empty());
739 assert!(pool.n_base <= crate::pipeline::FK_POOL_CAP);
741 }
742
743 #[test]
744 fn test_build_pipeline_config_only_entity_unknown_name() {
745 let schema = kyc_schema();
746 let err = build_pipeline_config(
747 "kyc",
748 1000,
749 42,
750 "medium",
751 0.1,
752 0.3,
753 &schema,
754 "kyc_test",
755 "parquet",
756 false,
757 "parquet",
758 Some("not_a_real_entity"),
759 )
760 .unwrap_err();
761 assert!(err.contains("not_a_real_entity"));
762 assert!(err.contains("natural_person"));
763 }
764
765 #[test]
766 fn test_deterministic_run_id_stable_and_sensitive() {
767 let a = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
768 let b = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
769 assert_eq!(a, b);
770 let c = deterministic_run_id("kyc", 1000, 43, "medium", 0.1, 0.3, "en", None, None);
771 assert_ne!(a, c);
772 }
773
774 #[test]
779 fn test_deterministic_run_id_sensitive_to_singleton_fraction_and_locale() {
780 let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
781 let diff_fraction =
782 deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.5, "en", None, None);
783 let diff_locale =
784 deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "fr", None, None);
785 assert_ne!(base, diff_fraction);
786 assert_ne!(base, diff_locale);
787 assert_ne!(diff_fraction, diff_locale);
788 }
789
790 #[test]
794 fn test_deterministic_run_id_sensitive_to_only_entity() {
795 let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
796 let only_a = deterministic_run_id(
797 "kyc",
798 1000,
799 42,
800 "medium",
801 0.1,
802 0.3,
803 "en",
804 Some("natural_person"),
805 None,
806 );
807 let only_b = deterministic_run_id(
808 "kyc",
809 1000,
810 42,
811 "medium",
812 0.1,
813 0.3,
814 "en",
815 Some("legal_entity"),
816 None,
817 );
818 assert_ne!(base, only_a);
819 assert_ne!(base, only_b);
820 assert_ne!(only_a, only_b);
821 }
822
823 #[test]
829 fn test_deterministic_run_id_sensitive_to_chunk_size() {
830 let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
831 let chunked_a =
832 deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, Some(500));
833 let chunked_b =
834 deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, Some(250));
835 assert_ne!(base, chunked_a);
836 assert_ne!(base, chunked_b);
837 assert_ne!(chunked_a, chunked_b);
838 }
839
840 #[test]
841 fn test_default_singleton_master_fraction_known_and_unknown() {
842 assert_eq!(default_singleton_master_fraction("light"), 0.50);
843 assert_eq!(default_singleton_master_fraction("hell"), 0.10);
844 assert_eq!(
846 default_singleton_master_fraction("bogus"),
847 default_singleton_master_fraction("medium")
848 );
849 }
850}