1use serde::Serialize;
7
8use crate::schema::{DomainSchema, build_pipeline_config};
9
10fn contains_any(s: &str, patterns: &[&str]) -> bool {
12 patterns.iter().any(|p| s.contains(p))
13}
14
15fn match_utility(col_name: &str, col_type: &str) -> f64 {
17 match col_type {
18 "boolean" => 0.0,
19 "int" | "float" => 0.2,
20 _ => {
21 let lower = col_name.to_lowercase();
22 if contains_any(&lower, &["record_id", "row_id"]) {
23 0.0
24 } else if lower.ends_with("_id") {
25 0.3
26 } else if contains_any(&lower, &["name", "email", "address", "phone", "ssn", "tax"]) {
27 1.0
28 } else if contains_any(
29 &lower,
30 &["date", "birth", "city", "country", "state", "postal"],
31 ) {
32 0.7
33 } else if contains_any(&lower, &["company", "legal", "trading", "registration"]) {
34 0.8
35 } else {
36 0.1
44 }
45 }
46 }
47}
48
49fn base_noise_damage(col_name: &str, col_type: &str) -> f64 {
51 let lower = col_name.to_lowercase();
52 match col_type {
53 "boolean" => 0.0,
54 "int" | "float" => 0.0,
55 "date" | "datetime" => {
56 if contains_any(&lower, &["birth", "dob"]) {
57 0.3
58 } else {
59 0.2
60 }
61 }
62 _ => {
63 if contains_any(&lower, &["email", "ssn", "phone", "mobile", "telephone"]) {
64 0.8
65 } else if contains_any(
66 &lower,
67 &[
68 "tax_id",
69 "registration",
70 "national_id",
71 "passport",
72 "account",
73 "barcode",
74 "pan",
75 "medicare",
76 ],
77 ) {
78 0.6
79 } else if contains_any(
80 &lower,
81 &[
82 "first_name",
83 "last_name",
84 "given_name",
85 "family_name",
86 "middle_name",
87 ],
88 ) {
89 0.4
90 } else if contains_any(
91 &lower,
92 &["address", "street", "city", "postal", "state", "country"],
93 ) {
94 0.5
95 } else if contains_any(&lower, &["date", "birth", "dob"]) {
96 0.3
97 } else {
98 0.4
101 }
102 }
103 }
104}
105
106fn parse_hn_id_fields(config_json: &str) -> Vec<String> {
109 #[derive(serde::Deserialize)]
110 struct HnConfigLight {
111 #[serde(default)]
112 pattern: String,
113 #[serde(default)]
114 id_fields: Vec<String>,
115 #[serde(default)]
116 mix_field: String,
117 #[serde(default)]
118 first_name_col: String,
119 #[serde(default)]
120 last_name_col: String,
121 #[serde(default)]
122 dob_col: String,
123 #[serde(default)]
124 email_col: String,
125 #[serde(default)]
126 ssn_col: String,
127 #[serde(default)]
128 phone_col: String,
129 #[serde(default)]
130 address_fields: Vec<String>,
131 }
132
133 let cfg: HnConfigLight = serde_json::from_str(config_json).unwrap_or(HnConfigLight {
134 pattern: String::new(),
135 id_fields: vec![],
136 mix_field: String::new(),
137 first_name_col: String::new(),
138 last_name_col: String::new(),
139 dob_col: String::new(),
140 email_col: String::new(),
141 ssn_col: String::new(),
142 phone_col: String::new(),
143 address_fields: vec![],
144 });
145
146 match cfg.pattern.as_str() {
147 "same_field" => cfg.id_fields,
148 "mix_identifier" => {
149 if cfg.mix_field.is_empty() {
150 vec![]
151 } else {
152 vec![cfg.mix_field]
153 }
154 }
155 "same_name_different_everything" => vec![cfg.first_name_col, cfg.last_name_col],
156 "same_email" => vec![cfg.email_col],
157 "same_ssn" => vec![cfg.ssn_col],
158 "same_phone" => vec![cfg.phone_col],
159 "same_address" => {
160 if cfg.address_fields.is_empty() {
161 vec!["address_line1".into(), "postal_code".into()]
162 } else {
163 cfg.address_fields
164 }
165 }
166 "same_name_dob" => vec![cfg.first_name_col, cfg.last_name_col, cfg.dob_col],
167 _ => vec![],
168 }
169 .into_iter()
170 .filter(|s| !s.is_empty())
171 .collect()
172}
173
174#[derive(Debug, Serialize)]
177pub struct ColReliability {
178 pub name: String,
179 pub col_type: String,
180 pub noise_damage: f64,
181 pub hn_risk: f64,
182 pub reliability: f64,
183}
184
185#[derive(Debug, Serialize)]
186pub struct EntityDifficulty {
187 pub name: String,
188 pub n_base: usize,
189 pub n_dup: usize,
190 pub true_pairs: usize,
191 pub hard_neg_pairs: usize,
192 pub guaranteed_fp: usize,
193 pub guaranteed_fn: usize,
194 pub columns: Vec<ColReliability>,
195}
196
197#[derive(Debug, Serialize)]
198pub struct DifficultyReport {
199 pub domain: String,
200 pub difficulty: String,
201 pub size: usize,
202 pub total_true_pairs: usize,
203 pub total_hard_neg_pairs: usize,
204 pub total_guaranteed_fp: usize,
205 pub total_guaranteed_fn: usize,
206 pub precision_max: f64,
207 pub recall_max: f64,
208 pub f1_max: f64,
209 pub entities: Vec<EntityDifficulty>,
210}
211
212#[derive(serde::Deserialize)]
215struct ColDefLight {
216 name: String,
217 #[serde(rename = "type", default = "default_str")]
218 col_type: String,
219}
220
221fn default_str() -> String {
222 "string".to_string()
223}
224
225pub fn estimate_difficulty(
226 domain: &str,
227 size: usize,
228 seed: u64,
229 difficulty: &str,
230 hard_neg_ratio: f64,
231 schema: &DomainSchema,
232) -> Result<DifficultyReport, String> {
233 let singleton_master_fraction = crate::schema::default_singleton_master_fraction(difficulty);
234 let run_id = crate::schema::deterministic_run_id(
238 domain,
239 size,
240 seed,
241 difficulty,
242 hard_neg_ratio,
243 singleton_master_fraction,
244 "en",
245 None,
246 None,
247 );
248 let config = build_pipeline_config(
249 domain,
250 size,
251 seed,
252 difficulty,
253 hard_neg_ratio,
254 singleton_master_fraction,
255 schema,
256 &run_id,
257 "ipc",
258 false,
259 "ipc",
260 None,
261 )?;
262
263 let mut hn_id_fields: std::collections::HashMap<String, Vec<String>> =
265 std::collections::HashMap::new();
266 for hn in &schema.hn_types {
267 let fields = parse_hn_id_fields(&hn.config_json);
268 hn_id_fields
269 .entry(hn.entity_type.clone())
270 .or_default()
271 .extend(fields);
272 }
273
274 let mut entities = Vec::new();
275 let mut total_true_pairs = 0usize;
276 let mut total_hard_neg_pairs = 0usize;
277 let mut total_guaranteed_fp = 0usize;
278 let mut total_guaranteed_fn = 0usize;
279 let mut total_guaranteed_fp_f = 0.0f64;
286 let mut total_guaranteed_fn_f = 0.0f64;
287
288 for plan in &config.entity_plans {
289 let poisoned: std::collections::HashSet<String> = hn_id_fields
290 .get(&plan.name)
291 .map(|v| v.iter().cloned().collect())
292 .unwrap_or_default();
293
294 let cols: Vec<ColDefLight> = serde_json::from_str(&plan.columns_json)
296 .map_err(|e| format!("parse columns for '{}': {}", plan.name, e))?;
297
298 let hn_pairs: usize = config
300 .hard_neg_types
301 .iter()
302 .filter(|h| h.entity_type == plan.name)
303 .map(|h| h.count)
304 .sum();
305
306 let n_dup: usize = plan.noise_types.iter().map(|n| n.count).sum();
308 let true_pairs = n_dup / 2;
309
310 let n_dup_f = n_dup.max(1) as f64;
330 let passes = config.noise_passes.max(1) as i32;
331 let mut col_reliability = Vec::new();
332 const BLOCKING_UTILITY_THRESHOLD: f64 = 0.7;
341 let mut blocking_fn_fail = 1.0f64;
342 let mut matching_fn_fail = 1.0f64;
343 let mut blocking_fp_fail = 1.0f64; let mut matching_fp_fail = 1.0f64;
345
346 for col in &cols {
347 let base_damage = base_noise_damage(&col.name, &col.col_type);
348 let is_dict_encoded =
365 crate::fast_template::resolve_dict_encoded_column(&col.name).is_some();
366 let p_single: f64 = if is_dict_encoded {
367 0.0
368 } else {
369 plan.noise_types
370 .iter()
371 .filter(|n| {
372 crate::pipeline::noise_type_targets_column(&n.noise_type, &col.name)
373 })
374 .map(|n| n.count as f64 / n_dup_f)
375 .sum::<f64>()
376 .min(1.0)
377 };
378 let p_touched = 1.0 - (1.0 - p_single).powi(passes);
379 let damage = base_damage * p_touched;
380 let util = match_utility(&col.name, &col.col_type);
381 let is_hn_id = poisoned.contains(&col.name);
382 let hn_risk = if is_hn_id { 1.0 } else { 0.0 };
383
384 let rel_fn = util * (1.0 - damage);
386
387 let rel_fp = util * (1.0 - hn_risk) * (1.0 - damage);
389
390 if util > 0.0 {
391 matching_fn_fail *= 1.0 - rel_fn;
392 matching_fp_fail *= 1.0 - rel_fp;
393 if util >= BLOCKING_UTILITY_THRESHOLD {
394 blocking_fn_fail *= 1.0 - rel_fn;
395 blocking_fp_fail *= 1.0 - rel_fp;
396 }
397 }
398
399 col_reliability.push(ColReliability {
400 name: col.name.clone(),
401 col_type: col.col_type.clone(),
402 noise_damage: damage,
403 hn_risk,
404 reliability: rel_fn.min(rel_fp),
405 });
406 }
407
408 let recall_blocking = if blocking_fn_fail < 1.0 {
412 1.0 - blocking_fn_fail
413 } else {
414 1.0
415 };
416 let precision_blocking = if blocking_fp_fail < 1.0 {
417 1.0 - blocking_fp_fail
418 } else {
419 1.0
420 };
421 let recall_combined = recall_blocking * (1.0 - matching_fn_fail);
422 let precision_combined = precision_blocking * (1.0 - matching_fp_fail);
423
424 let guaranteed_fp_f = if hn_pairs > 0 {
425 hn_pairs as f64 * (1.0 - precision_combined)
427 } else {
428 0.0
429 };
430 let guaranteed_fn_f = true_pairs as f64 * (1.0 - recall_combined);
431 let guaranteed_fp = guaranteed_fp_f as usize;
432 let guaranteed_fn = guaranteed_fn_f as usize;
433
434 total_true_pairs += true_pairs;
435 total_hard_neg_pairs += hn_pairs;
436 total_guaranteed_fp += guaranteed_fp;
437 total_guaranteed_fn += guaranteed_fn;
438 total_guaranteed_fp_f += guaranteed_fp_f;
439 total_guaranteed_fn_f += guaranteed_fn_f;
440
441 entities.push(EntityDifficulty {
442 name: plan.name.clone(),
443 n_base: plan.n_base,
444 n_dup,
445 true_pairs,
446 hard_neg_pairs: hn_pairs,
447 guaranteed_fp,
448 guaranteed_fn,
449 columns: col_reliability,
450 });
451 }
452
453 let tp = total_true_pairs.max(1) as f64;
454 let fp = total_guaranteed_fp_f;
455 let fn_ = total_guaranteed_fn_f;
456
457 let precision_max = tp / (tp + fp);
458 let recall_max = tp / (tp + fn_);
459 let f1_max = 2.0 * precision_max * recall_max / (precision_max + recall_max);
460
461 Ok(DifficultyReport {
462 domain: domain.to_string(),
463 difficulty: difficulty.to_string(),
464 size,
465 total_true_pairs,
466 total_hard_neg_pairs,
467 total_guaranteed_fp,
468 total_guaranteed_fn,
469 precision_max,
470 recall_max,
471 f1_max,
472 entities,
473 })
474}
475
476#[cfg(test)]
477mod tests {
478 use super::*;
479 use crate::schema::load_schema;
480
481 fn kyc_schema() -> DomainSchema {
482 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
483 load_schema("kyc", &dir).expect("load kyc.json")
484 }
485
486 #[test]
487 fn test_estimate_difficulty_basic_shape() {
488 let schema = kyc_schema();
489 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
490 assert_eq!(report.domain, "kyc");
491 assert_eq!(report.difficulty, "medium");
492 assert_eq!(report.entities.len(), 2);
493 for e in &report.entities {
494 assert!(!e.columns.is_empty());
495 }
496 }
497
498 #[test]
499 fn test_estimate_difficulty_f1_bounds() {
500 let schema = kyc_schema();
501 let report = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
502 assert!(report.precision_max > 0.0 && report.precision_max <= 1.0);
503 assert!(report.recall_max > 0.0 && report.recall_max <= 1.0);
504 assert!(report.f1_max > 0.0 && report.f1_max <= 1.0);
505 }
506
507 #[test]
508 fn test_estimate_difficulty_hell_harder_than_light() {
509 let schema = kyc_schema();
512 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
513 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
514 assert!(hell.f1_max <= light.f1_max);
515 }
516
517 #[test]
518 fn test_estimate_difficulty_tiers_ordered() {
519 let schema = kyc_schema();
523 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
524 let medium = estimate_difficulty("kyc", 5000, 42, "medium", 0.1, &schema).unwrap();
525 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
526 assert!(medium.f1_max <= light.f1_max);
527 assert!(hell.f1_max <= medium.f1_max);
528 }
529
530 #[test]
531 fn test_estimate_difficulty_tiers_ordered_healthcare() {
532 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
544 let schema = load_schema("healthcare", &dir).expect("load healthcare.json");
545 let light = estimate_difficulty("healthcare", 5000, 42, "light", 0.1, &schema).unwrap();
546 let medium = estimate_difficulty("healthcare", 5000, 42, "medium", 0.1, &schema).unwrap();
547 let hell = estimate_difficulty("healthcare", 5000, 42, "hell", 0.1, &schema).unwrap();
548 assert!(medium.f1_max <= light.f1_max);
549 assert!(hell.f1_max <= medium.f1_max);
550 }
551
552 #[test]
553 fn test_estimate_difficulty_tiers_ordered_aviation() {
554 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
567 let schema = load_schema("aviation", &dir).expect("load aviation.json");
568 let light = estimate_difficulty("aviation", 5000, 42, "light", 0.1, &schema).unwrap();
569 let medium = estimate_difficulty("aviation", 5000, 42, "medium", 0.1, &schema).unwrap();
570 let hell = estimate_difficulty("aviation", 5000, 42, "hell", 0.1, &schema).unwrap();
571 assert!(medium.f1_max <= light.f1_max);
572 assert!(hell.f1_max <= medium.f1_max);
573 }
574
575 #[test]
576 fn test_estimate_difficulty_tiers_ordered_all_domains() {
577 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
587 let eps = 1e-6;
588 let mut violations = Vec::new();
589 for entry in std::fs::read_dir(&dir).expect("read schemas dir") {
590 let path = entry.expect("dir entry").path();
591 if path.extension().and_then(|e| e.to_str()) != Some("json") {
592 continue;
593 }
594 let domain = path.file_stem().unwrap().to_str().unwrap().to_string();
595 let schema = load_schema(&domain, &dir).expect("load schema");
596 let light = estimate_difficulty(&domain, 5000, 42, "light", 0.1, &schema).unwrap();
597 let medium = estimate_difficulty(&domain, 5000, 42, "medium", 0.1, &schema).unwrap();
598 let hell = estimate_difficulty(&domain, 5000, 42, "hell", 0.1, &schema).unwrap();
599 if medium.f1_max > light.f1_max + eps {
600 violations.push(format!(
601 "{domain}: medium ({:.4}) > light ({:.4})",
602 medium.f1_max, light.f1_max
603 ));
604 }
605 if hell.f1_max > medium.f1_max + eps {
606 violations.push(format!(
607 "{domain}: hell ({:.4}) > medium ({:.4})",
608 hell.f1_max, medium.f1_max
609 ));
610 }
611 }
612 assert!(
613 violations.is_empty(),
614 "monotonicity violations:\n{}",
615 violations.join("\n")
616 );
617 }
618
619 #[test]
620 fn test_estimate_difficulty_deterministic() {
621 let schema = kyc_schema();
622 let a = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
623 let b = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
624 assert_eq!(a.total_true_pairs, b.total_true_pairs);
625 assert_eq!(a.f1_max, b.f1_max);
626 }
627
628 #[test]
629 fn test_estimate_difficulty_zero_hard_neg_ratio_no_fp() {
630 let schema = kyc_schema();
631 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.0, &schema).unwrap();
632 assert_eq!(report.total_hard_neg_pairs, 0);
633 assert_eq!(report.total_guaranteed_fp, 0);
634 }
635}