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