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 p_single: f64 = plan
349 .noise_types
350 .iter()
351 .filter(|n| crate::pipeline::noise_type_targets_column(&n.noise_type, &col.name))
352 .map(|n| n.count as f64 / n_dup_f)
353 .sum::<f64>()
354 .min(1.0);
355 let p_touched = 1.0 - (1.0 - p_single).powi(passes);
356 let damage = base_damage * p_touched;
357 let util = match_utility(&col.name, &col.col_type);
358 let is_hn_id = poisoned.contains(&col.name);
359 let hn_risk = if is_hn_id { 1.0 } else { 0.0 };
360
361 let rel_fn = util * (1.0 - damage);
363
364 let rel_fp = util * (1.0 - hn_risk) * (1.0 - damage);
366
367 if util > 0.0 {
368 matching_fn_fail *= 1.0 - rel_fn;
369 matching_fp_fail *= 1.0 - rel_fp;
370 if util >= BLOCKING_UTILITY_THRESHOLD {
371 blocking_fn_fail *= 1.0 - rel_fn;
372 blocking_fp_fail *= 1.0 - rel_fp;
373 }
374 }
375
376 col_reliability.push(ColReliability {
377 name: col.name.clone(),
378 col_type: col.col_type.clone(),
379 noise_damage: damage,
380 hn_risk,
381 reliability: rel_fn.min(rel_fp),
382 });
383 }
384
385 let recall_blocking = if blocking_fn_fail < 1.0 {
389 1.0 - blocking_fn_fail
390 } else {
391 1.0
392 };
393 let precision_blocking = if blocking_fp_fail < 1.0 {
394 1.0 - blocking_fp_fail
395 } else {
396 1.0
397 };
398 let recall_combined = recall_blocking * (1.0 - matching_fn_fail);
399 let precision_combined = precision_blocking * (1.0 - matching_fp_fail);
400
401 let guaranteed_fp_f = if hn_pairs > 0 {
402 hn_pairs as f64 * (1.0 - precision_combined)
404 } else {
405 0.0
406 };
407 let guaranteed_fn_f = true_pairs as f64 * (1.0 - recall_combined);
408 let guaranteed_fp = guaranteed_fp_f as usize;
409 let guaranteed_fn = guaranteed_fn_f as usize;
410
411 total_true_pairs += true_pairs;
412 total_hard_neg_pairs += hn_pairs;
413 total_guaranteed_fp += guaranteed_fp;
414 total_guaranteed_fn += guaranteed_fn;
415 total_guaranteed_fp_f += guaranteed_fp_f;
416 total_guaranteed_fn_f += guaranteed_fn_f;
417
418 entities.push(EntityDifficulty {
419 name: plan.name.clone(),
420 n_base: plan.n_base,
421 n_dup,
422 true_pairs,
423 hard_neg_pairs: hn_pairs,
424 guaranteed_fp,
425 guaranteed_fn,
426 columns: col_reliability,
427 });
428 }
429
430 let tp = total_true_pairs.max(1) as f64;
431 let fp = total_guaranteed_fp_f;
432 let fn_ = total_guaranteed_fn_f;
433
434 let precision_max = tp / (tp + fp);
435 let recall_max = tp / (tp + fn_);
436 let f1_max = 2.0 * precision_max * recall_max / (precision_max + recall_max);
437
438 Ok(DifficultyReport {
439 domain: domain.to_string(),
440 difficulty: difficulty.to_string(),
441 size,
442 total_true_pairs,
443 total_hard_neg_pairs,
444 total_guaranteed_fp,
445 total_guaranteed_fn,
446 precision_max,
447 recall_max,
448 f1_max,
449 entities,
450 })
451}
452
453#[cfg(test)]
454mod tests {
455 use super::*;
456 use crate::schema::load_schema;
457
458 fn kyc_schema() -> DomainSchema {
459 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
460 load_schema("kyc", &dir).expect("load kyc.json")
461 }
462
463 #[test]
464 fn test_estimate_difficulty_basic_shape() {
465 let schema = kyc_schema();
466 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
467 assert_eq!(report.domain, "kyc");
468 assert_eq!(report.difficulty, "medium");
469 assert_eq!(report.entities.len(), 2);
470 for e in &report.entities {
471 assert!(!e.columns.is_empty());
472 }
473 }
474
475 #[test]
476 fn test_estimate_difficulty_f1_bounds() {
477 let schema = kyc_schema();
478 let report = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
479 assert!(report.precision_max > 0.0 && report.precision_max <= 1.0);
480 assert!(report.recall_max > 0.0 && report.recall_max <= 1.0);
481 assert!(report.f1_max > 0.0 && report.f1_max <= 1.0);
482 }
483
484 #[test]
485 fn test_estimate_difficulty_hell_harder_than_light() {
486 let schema = kyc_schema();
489 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
490 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
491 assert!(hell.f1_max <= light.f1_max);
492 }
493
494 #[test]
495 fn test_estimate_difficulty_tiers_ordered() {
496 let schema = kyc_schema();
500 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
501 let medium = estimate_difficulty("kyc", 5000, 42, "medium", 0.1, &schema).unwrap();
502 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
503 assert!(medium.f1_max <= light.f1_max);
504 assert!(hell.f1_max <= medium.f1_max);
505 }
506
507 #[test]
508 fn test_estimate_difficulty_tiers_ordered_healthcare() {
509 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
521 let schema = load_schema("healthcare", &dir).expect("load healthcare.json");
522 let light = estimate_difficulty("healthcare", 5000, 42, "light", 0.1, &schema).unwrap();
523 let medium = estimate_difficulty("healthcare", 5000, 42, "medium", 0.1, &schema).unwrap();
524 let hell = estimate_difficulty("healthcare", 5000, 42, "hell", 0.1, &schema).unwrap();
525 assert!(medium.f1_max <= light.f1_max);
526 assert!(hell.f1_max <= medium.f1_max);
527 }
528
529 #[test]
530 fn test_estimate_difficulty_tiers_ordered_aviation() {
531 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
544 let schema = load_schema("aviation", &dir).expect("load aviation.json");
545 let light = estimate_difficulty("aviation", 5000, 42, "light", 0.1, &schema).unwrap();
546 let medium = estimate_difficulty("aviation", 5000, 42, "medium", 0.1, &schema).unwrap();
547 let hell = estimate_difficulty("aviation", 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_all_domains() {
554 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
564 let eps = 1e-6;
565 let mut violations = Vec::new();
566 for entry in std::fs::read_dir(&dir).expect("read schemas dir") {
567 let path = entry.expect("dir entry").path();
568 if path.extension().and_then(|e| e.to_str()) != Some("json") {
569 continue;
570 }
571 let domain = path.file_stem().unwrap().to_str().unwrap().to_string();
572 let schema = load_schema(&domain, &dir).expect("load schema");
573 let light = estimate_difficulty(&domain, 5000, 42, "light", 0.1, &schema).unwrap();
574 let medium = estimate_difficulty(&domain, 5000, 42, "medium", 0.1, &schema).unwrap();
575 let hell = estimate_difficulty(&domain, 5000, 42, "hell", 0.1, &schema).unwrap();
576 if medium.f1_max > light.f1_max + eps {
577 violations.push(format!(
578 "{domain}: medium ({:.4}) > light ({:.4})",
579 medium.f1_max, light.f1_max
580 ));
581 }
582 if hell.f1_max > medium.f1_max + eps {
583 violations.push(format!(
584 "{domain}: hell ({:.4}) > medium ({:.4})",
585 hell.f1_max, medium.f1_max
586 ));
587 }
588 }
589 assert!(
590 violations.is_empty(),
591 "monotonicity violations:\n{}",
592 violations.join("\n")
593 );
594 }
595
596 #[test]
597 fn test_estimate_difficulty_deterministic() {
598 let schema = kyc_schema();
599 let a = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
600 let b = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
601 assert_eq!(a.total_true_pairs, b.total_true_pairs);
602 assert_eq!(a.f1_max, b.f1_max);
603 }
604
605 #[test]
606 fn test_estimate_difficulty_zero_hard_neg_ratio_no_fp() {
607 let schema = kyc_schema();
608 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.0, &schema).unwrap();
609 assert_eq!(report.total_hard_neg_pairs, 0);
610 assert_eq!(report.total_guaranteed_fp, 0);
611 }
612}