1use serde::Serialize;
8
9use crate::schema::{DomainSchema, build_pipeline_config};
10
11fn contains_any(s: &str, patterns: &[&str]) -> bool {
13 patterns.iter().any(|p| s.contains(p))
14}
15
16fn match_utility(col_name: &str, col_type: &str) -> f64 {
18 match col_type {
19 "boolean" => 0.0,
20 "int" | "float" => 0.2,
21 _ => {
22 let lower = col_name.to_lowercase();
23 if contains_any(&lower, &["record_id", "row_id"]) {
24 0.0
25 } else if lower.ends_with("_id") {
26 0.3
27 } else if contains_any(&lower, &["name", "email", "address", "phone", "ssn", "tax"]) {
28 1.0
29 } else if contains_any(
30 &lower,
31 &["date", "birth", "city", "country", "state", "postal"],
32 ) {
33 0.7
34 } else if contains_any(&lower, &["company", "legal", "trading", "registration"]) {
35 0.8
36 } else {
37 0.5
38 }
39 }
40 }
41}
42
43fn base_noise_damage(col_name: &str, col_type: &str) -> f64 {
45 let lower = col_name.to_lowercase();
46 match col_type {
47 "boolean" => 0.0,
48 "int" | "float" => 0.0,
49 "date" | "datetime" => {
50 if contains_any(&lower, &["birth", "dob"]) {
51 0.3
52 } else {
53 0.2
54 }
55 }
56 _ => {
57 if contains_any(&lower, &["email", "ssn", "phone", "mobile", "telephone"]) {
58 0.8
59 } else if contains_any(
60 &lower,
61 &[
62 "tax_id",
63 "registration",
64 "national_id",
65 "passport",
66 "account",
67 "barcode",
68 "pan",
69 "medicare",
70 ],
71 ) {
72 0.6
73 } else if contains_any(
74 &lower,
75 &[
76 "first_name",
77 "last_name",
78 "given_name",
79 "family_name",
80 "middle_name",
81 ],
82 ) {
83 0.4
84 } else if contains_any(
85 &lower,
86 &["address", "street", "city", "postal", "state", "country"],
87 ) {
88 0.5
89 } else if contains_any(&lower, &["date", "birth", "dob"]) {
90 0.3
91 } else {
92 0.4
95 }
96 }
97 }
98}
99
100fn parse_hn_id_fields(config_json: &str) -> Vec<String> {
103 #[derive(serde::Deserialize)]
104 struct HnConfigLight {
105 #[serde(default)]
106 pattern: String,
107 #[serde(default)]
108 id_fields: Vec<String>,
109 #[serde(default)]
110 mix_field: String,
111 #[serde(default)]
112 first_name_col: String,
113 #[serde(default)]
114 last_name_col: String,
115 #[serde(default)]
116 dob_col: String,
117 #[serde(default)]
118 email_col: String,
119 #[serde(default)]
120 ssn_col: String,
121 #[serde(default)]
122 phone_col: String,
123 #[serde(default)]
124 address_fields: Vec<String>,
125 }
126
127 let cfg: HnConfigLight = serde_json::from_str(config_json).unwrap_or(HnConfigLight {
128 pattern: String::new(),
129 id_fields: vec![],
130 mix_field: String::new(),
131 first_name_col: String::new(),
132 last_name_col: String::new(),
133 dob_col: String::new(),
134 email_col: String::new(),
135 ssn_col: String::new(),
136 phone_col: String::new(),
137 address_fields: vec![],
138 });
139
140 match cfg.pattern.as_str() {
141 "same_field" => cfg.id_fields,
142 "mix_identifier" => {
143 if cfg.mix_field.is_empty() {
144 vec![]
145 } else {
146 vec![cfg.mix_field]
147 }
148 }
149 "same_name_different_everything" => vec![cfg.first_name_col, cfg.last_name_col],
150 "same_email" => vec![cfg.email_col],
151 "same_ssn" => vec![cfg.ssn_col],
152 "same_phone" => vec![cfg.phone_col],
153 "same_address" => {
154 if cfg.address_fields.is_empty() {
155 vec!["address_line1".into(), "postal_code".into()]
156 } else {
157 cfg.address_fields
158 }
159 }
160 "same_name_dob" => vec![cfg.first_name_col, cfg.last_name_col, cfg.dob_col],
161 _ => vec![],
162 }
163 .into_iter()
164 .filter(|s| !s.is_empty())
165 .collect()
166}
167
168#[derive(Debug, Serialize)]
171pub struct ColReliability {
172 pub name: String,
173 pub col_type: String,
174 pub noise_damage: f64,
175 pub hn_risk: f64,
176 pub reliability: f64,
177}
178
179#[derive(Debug, Serialize)]
180pub struct EntityDifficulty {
181 pub name: String,
182 pub n_base: usize,
183 pub n_dup: usize,
184 pub true_pairs: usize,
185 pub hard_neg_pairs: usize,
186 pub guaranteed_fp: usize,
187 pub guaranteed_fn: usize,
188 pub columns: Vec<ColReliability>,
189}
190
191#[derive(Debug, Serialize)]
192pub struct DifficultyReport {
193 pub domain: String,
194 pub difficulty: String,
195 pub size: usize,
196 pub total_true_pairs: usize,
197 pub total_hard_neg_pairs: usize,
198 pub total_guaranteed_fp: usize,
199 pub total_guaranteed_fn: usize,
200 pub precision_max: f64,
201 pub recall_max: f64,
202 pub f1_max: f64,
203 pub entities: Vec<EntityDifficulty>,
204}
205
206#[derive(serde::Deserialize)]
209struct ColDefLight {
210 name: String,
211 #[serde(rename = "type", default = "default_str")]
212 col_type: String,
213}
214
215fn default_str() -> String {
216 "string".to_string()
217}
218
219pub fn estimate_difficulty(
220 domain: &str,
221 size: usize,
222 seed: u64,
223 difficulty: &str,
224 hard_neg_ratio: f64,
225 schema: &DomainSchema,
226) -> Result<DifficultyReport, String> {
227 let singleton_master_fraction = crate::schema::default_singleton_master_fraction(difficulty);
228 let run_id = crate::schema::deterministic_run_id(
232 domain,
233 size,
234 seed,
235 difficulty,
236 hard_neg_ratio,
237 singleton_master_fraction,
238 "en",
239 );
240 let config = build_pipeline_config(
241 domain,
242 size,
243 seed,
244 difficulty,
245 hard_neg_ratio,
246 singleton_master_fraction,
247 schema,
248 &run_id,
249 "ipc",
250 false,
251 "ipc",
252 )?;
253
254 let mut hn_id_fields: std::collections::HashMap<String, Vec<String>> =
256 std::collections::HashMap::new();
257 for hn in &schema.hn_types {
258 let fields = parse_hn_id_fields(&hn.config_json);
259 hn_id_fields
260 .entry(hn.entity_type.clone())
261 .or_default()
262 .extend(fields);
263 }
264
265 let mut entities = Vec::new();
266 let mut total_true_pairs = 0usize;
267 let mut total_hard_neg_pairs = 0usize;
268 let mut total_guaranteed_fp = 0usize;
269 let mut total_guaranteed_fn = 0usize;
270
271 for plan in &config.entity_plans {
272 let poisoned: std::collections::HashSet<String> = hn_id_fields
273 .get(&plan.name)
274 .map(|v| v.iter().cloned().collect())
275 .unwrap_or_default();
276
277 let cols: Vec<ColDefLight> = serde_json::from_str(&plan.columns_json)
279 .map_err(|e| format!("parse columns for '{}': {}", plan.name, e))?;
280
281 let hn_pairs: usize = config
283 .hard_neg_types
284 .iter()
285 .filter(|h| h.entity_type == plan.name)
286 .map(|h| h.count)
287 .sum();
288
289 let n_dup: usize = plan.noise_types.iter().map(|n| n.count).sum();
291 let true_pairs = n_dup / 2;
292
293 let n_dup_f = n_dup.max(1) as f64;
313 let passes = config.noise_passes.max(1) as i32;
314 let mut col_reliability = Vec::new();
315 let mut best_fn_reliability = 0.0f64;
316 let mut best_fp_reliability = 0.0f64; for col in &cols {
319 let base_damage = base_noise_damage(&col.name, &col.col_type);
320 let p_single: f64 = plan
321 .noise_types
322 .iter()
323 .filter(|n| crate::pipeline::noise_type_targets_column(&n.noise_type, &col.name))
324 .map(|n| n.count as f64 / n_dup_f)
325 .sum::<f64>()
326 .min(1.0);
327 let p_touched = 1.0 - (1.0 - p_single).powi(passes);
328 let damage = base_damage * p_touched;
329 let util = match_utility(&col.name, &col.col_type);
330 let is_hn_id = poisoned.contains(&col.name);
331 let hn_risk = if is_hn_id { 1.0 } else { 0.0 };
332
333 let rel_fn = util * (1.0 - damage);
335
336 let rel_fp = util * (1.0 - hn_risk) * (1.0 - damage);
338
339 if rel_fn > best_fn_reliability {
340 best_fn_reliability = rel_fn;
341 }
342 if rel_fp > best_fp_reliability {
343 best_fp_reliability = rel_fp;
344 }
345
346 col_reliability.push(ColReliability {
347 name: col.name.clone(),
348 col_type: col.col_type.clone(),
349 noise_damage: damage,
350 hn_risk,
351 reliability: rel_fn.min(rel_fp),
352 });
353 }
354
355 let guaranteed_fp = if hn_pairs > 0 {
356 (hn_pairs as f64 * (1.0 - best_fp_reliability)) as usize
358 } else {
359 0
360 };
361
362 let guaranteed_fn = (true_pairs as f64 * (1.0 - best_fn_reliability)) as usize;
363
364 total_true_pairs += true_pairs;
365 total_hard_neg_pairs += hn_pairs;
366 total_guaranteed_fp += guaranteed_fp;
367 total_guaranteed_fn += guaranteed_fn;
368
369 entities.push(EntityDifficulty {
370 name: plan.name.clone(),
371 n_base: plan.n_base,
372 n_dup,
373 true_pairs,
374 hard_neg_pairs: hn_pairs,
375 guaranteed_fp,
376 guaranteed_fn,
377 columns: col_reliability,
378 });
379 }
380
381 let tp = total_true_pairs.max(1) as f64;
382 let fp = total_guaranteed_fp as f64;
383 let fn_ = total_guaranteed_fn as f64;
384
385 let precision_max = tp / (tp + fp);
386 let recall_max = tp / (tp + fn_);
387 let f1_max = 2.0 * precision_max * recall_max / (precision_max + recall_max);
388
389 Ok(DifficultyReport {
390 domain: domain.to_string(),
391 difficulty: difficulty.to_string(),
392 size,
393 total_true_pairs,
394 total_hard_neg_pairs,
395 total_guaranteed_fp,
396 total_guaranteed_fn,
397 precision_max,
398 recall_max,
399 f1_max,
400 entities,
401 })
402}
403
404#[cfg(test)]
405mod tests {
406 use super::*;
407 use crate::schema::load_schema;
408
409 fn kyc_schema() -> DomainSchema {
410 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
411 load_schema("kyc", &dir).expect("load kyc.json")
412 }
413
414 #[test]
415 fn test_estimate_difficulty_basic_shape() {
416 let schema = kyc_schema();
417 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
418 assert_eq!(report.domain, "kyc");
419 assert_eq!(report.difficulty, "medium");
420 assert_eq!(report.entities.len(), 2);
421 for e in &report.entities {
422 assert!(!e.columns.is_empty());
423 }
424 }
425
426 #[test]
427 fn test_estimate_difficulty_f1_bounds() {
428 let schema = kyc_schema();
429 let report = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
430 assert!(report.precision_max > 0.0 && report.precision_max <= 1.0);
431 assert!(report.recall_max > 0.0 && report.recall_max <= 1.0);
432 assert!(report.f1_max > 0.0 && report.f1_max <= 1.0);
433 }
434
435 #[test]
436 fn test_estimate_difficulty_hell_harder_than_light() {
437 let schema = kyc_schema();
440 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
441 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
442 assert!(hell.f1_max <= light.f1_max);
443 }
444
445 #[test]
446 fn test_estimate_difficulty_tiers_ordered() {
447 let schema = kyc_schema();
451 let light = estimate_difficulty("kyc", 5000, 42, "light", 0.1, &schema).unwrap();
452 let medium = estimate_difficulty("kyc", 5000, 42, "medium", 0.1, &schema).unwrap();
453 let hell = estimate_difficulty("kyc", 5000, 42, "hell", 0.1, &schema).unwrap();
454 assert!(medium.f1_max <= light.f1_max);
455 assert!(hell.f1_max <= medium.f1_max);
456 }
457
458 #[test]
459 fn test_estimate_difficulty_tiers_ordered_healthcare() {
460 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
472 let schema = load_schema("healthcare", &dir).expect("load healthcare.json");
473 let light = estimate_difficulty("healthcare", 5000, 42, "light", 0.1, &schema).unwrap();
474 let medium = estimate_difficulty("healthcare", 5000, 42, "medium", 0.1, &schema).unwrap();
475 let hell = estimate_difficulty("healthcare", 5000, 42, "hell", 0.1, &schema).unwrap();
476 assert!(medium.f1_max <= light.f1_max);
477 assert!(hell.f1_max <= medium.f1_max);
478 }
479
480 #[test]
481 fn test_estimate_difficulty_tiers_ordered_aviation() {
482 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
495 let schema = load_schema("aviation", &dir).expect("load aviation.json");
496 let light = estimate_difficulty("aviation", 5000, 42, "light", 0.1, &schema).unwrap();
497 let medium = estimate_difficulty("aviation", 5000, 42, "medium", 0.1, &schema).unwrap();
498 let hell = estimate_difficulty("aviation", 5000, 42, "hell", 0.1, &schema).unwrap();
499 assert!(medium.f1_max <= light.f1_max);
500 assert!(hell.f1_max <= medium.f1_max);
501 }
502
503 #[test]
504 fn test_estimate_difficulty_tiers_ordered_all_domains() {
505 let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas");
515 let eps = 1e-6;
516 let mut violations = Vec::new();
517 for entry in std::fs::read_dir(&dir).expect("read schemas dir") {
518 let path = entry.expect("dir entry").path();
519 if path.extension().and_then(|e| e.to_str()) != Some("json") {
520 continue;
521 }
522 let domain = path.file_stem().unwrap().to_str().unwrap().to_string();
523 let schema = load_schema(&domain, &dir).expect("load schema");
524 let light = estimate_difficulty(&domain, 5000, 42, "light", 0.1, &schema).unwrap();
525 let medium = estimate_difficulty(&domain, 5000, 42, "medium", 0.1, &schema).unwrap();
526 let hell = estimate_difficulty(&domain, 5000, 42, "hell", 0.1, &schema).unwrap();
527 if medium.f1_max > light.f1_max + eps {
528 violations.push(format!(
529 "{domain}: medium ({:.4}) > light ({:.4})",
530 medium.f1_max, light.f1_max
531 ));
532 }
533 if hell.f1_max > medium.f1_max + eps {
534 violations.push(format!(
535 "{domain}: hell ({:.4}) > medium ({:.4})",
536 hell.f1_max, medium.f1_max
537 ));
538 }
539 }
540 assert!(
541 violations.is_empty(),
542 "monotonicity violations:\n{}",
543 violations.join("\n")
544 );
545 }
546
547 #[test]
548 fn test_estimate_difficulty_deterministic() {
549 let schema = kyc_schema();
550 let a = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
551 let b = estimate_difficulty("kyc", 1000, 42, "medium", 0.1, &schema).unwrap();
552 assert_eq!(a.total_true_pairs, b.total_true_pairs);
553 assert_eq!(a.f1_max, b.f1_max);
554 }
555
556 #[test]
557 fn test_estimate_difficulty_zero_hard_neg_ratio_no_fp() {
558 let schema = kyc_schema();
559 let report = estimate_difficulty("kyc", 1000, 42, "medium", 0.0, &schema).unwrap();
560 assert_eq!(report.total_hard_neg_pairs, 0);
561 assert_eq!(report.total_guaranteed_fp, 0);
562 }
563}