Skip to main content

dupehell_core/
schema.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// EDUCATIONAL AND RESEARCH PURPOSES ONLY -- see ETHICS.md for prohibited uses.
5// No liability for misuse.
6
7use 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
22/// Infer an entity's identifier column: prefer `{entity_name}_id`, then
23/// `id`, then the first column ending in `_id`.
24fn 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    /// Relative population size of this entity within the domain, used to
50    /// split `total_unique` identities across entities (`build_pipeline_config`).
51    /// Defaults to `1.0` (every entity gets an equal share) when a schema
52    /// doesn't set it — every schema written before this field existed keeps
53    /// behaving exactly as before. Set explicit weights when entities aren't
54    /// realistically equal in population (e.g. aviation: far more
55    /// `passenger` identities than `airline` ones) — see `docs/` or ask
56    /// before picking numbers for a domain you don't know well; a wrong
57    /// weight is as misleading as the uniform default, just in the other
58    /// direction.
59    #[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 active for this difficulty, weighted equally. Curated
74    /// explicitly per level (rather than derived from a truncated,
75    /// order-dependent countdown) so that "broad" noise types — `typo` and
76    /// `visual`, which both corrupt name+address+phone+company on the same
77    /// record at once (see `pipeline::noise_type_targets_column`) — never
78    /// end up a *larger* share of the mix at a lower difficulty than at a
79    /// higher one. A prior countdown-weight scheme caused exactly that:
80    /// medium (4 active types) put 70% of its noise budget on typo+visual,
81    /// vs. only ~42% for hell (8 active types), making medium duplicates
82    /// *more* likely to have every strong matching field wiped out at once
83    /// than hell duplicates, despite medium being meant as the easier tier.
84    noise_types: &'static [&'static str],
85    /// Number of *independent* noise passes applied to each duplicate copy
86    /// (see `pipeline::run_pipeline`'s dup-generation loop). Each additional
87    /// pass draws its own noise_type from `noise_types` and applies it on
88    /// top of the previous pass's result.
89    ///
90    /// This is the actual difficulty lever between tiers with different
91    /// `noise_types.len()` — not the category list itself. Each entry's
92    /// per-pass weight is `1 / noise_types.len()` (`build_pipeline_config`
93    /// below), so adding categories to make a tier "harder" *dilutes* every
94    /// existing category's weight, including whichever one happens to guard
95    /// a domain's single most reliable linkage column — a tier with more
96    /// categories active can paradoxically end up *easier* on some schemas
97    /// (measured: a "hard" tier here, folded into hell, kept coming out
98    /// easier than "hell" on several domains despite having fewer
99    /// categories; then a hand-picked category weighted twice fixed 2 of 6
100    /// domains tested but not the rest, since the dominant reliable column
101    /// isn't always the same category across schemas). Passes sidestep this
102    /// entirely: the probability a given column is touched at least once
103    /// across `P` independent passes is `1 - (1 - p)^P` where `p` is its
104    /// single-pass weight — strictly increasing in `P` regardless of how
105    /// small `p` is, so more passes can never make a tier easier the way
106    /// more categories can. `difficulty::estimate_difficulty` models this
107    /// exact same formula.
108    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            // 3 passes: a column protected by exactly one of medium's 4
128            // equally-weighted categories (weight 1/4 -- e.g. a pure date
129            // column, only matched by "dates", not by "typo"'s broader
130            // predicate) is diluted well below light's worst case (2
131            // categories, weight 1/2). At passes=1 this let medium come out
132            // *easier* than light on 19/40 domains (found via
133            // scripts/validate_all_domains.py). 1 - (1 - 1/4)^3 ≈ 0.578 >
134            // light's single-pass 0.5, closing the gap for any category
135            // medium shares with light regardless of which column ends up
136            // being a given domain's escape hatch. Re-validated across all
137            // 40 domains with the same script after this change.
138            passes: 3,
139        },
140    ),
141    (
142        "hell",
143        DifficultySettings {
144            singleton: 0.10,
145            doublet: 0.20,
146            // `typo_aggressive` corrupts more of each string than plain
147            // `typo` (same target columns, more damage per hit — see
148            // `noise::typos`), and `unicode_pollution` gets its own bucket
149            // (rather than only ever showing up as one of `visual`'s five
150            // random sub-choices).
151            noise_types: &[
152                "typo_aggressive",
153                "visual",
154                "unicode_pollution",
155                "names",
156                "dates",
157                "identifiers",
158                "addresses",
159                "companies",
160                "extra",
161            ],
162            // 8 passes: medium was bumped from 1 to 3 passes to fix its own
163            // dilution against light (see the comment on medium's `passes`
164            // above), which raised medium's worst-case single-category
165            // touch probability to 1-(1-1/4)^3 ≈ 0.578. With 9
166            // equally-weighted categories (weight 1/9 each), hell needs
167            // 1-(1-1/9)^passes >= 0.578 to stay at least as hard as medium
168            // on any column protected by only one category -- solving gives
169            // passes >= 7.33, rounded up to 8 (1-(1-1/9)^8 ≈ 0.617).
170            // Re-validated across all 40 domains via
171            // scripts/validate_all_domains.py (0 crashes, 0 monotonicity
172            // violations) after this change.
173            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
190/// Generate a domain-unique run ID based on the current Unix timestamp (hex).
191pub 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/// Deterministic run ID derived from generation parameters, so the same
199/// (domain, size, seed, difficulty, hard_neg_ratio, singleton_master_fraction,
200/// locale) always produces the same output filename regardless of output
201/// format (IPC vs Parquet) or how many times it's run — and, just as
202/// important, a *different* filename whenever any of these parameters
203/// differ, since they all affect the generated data. `singleton_master_fraction`
204/// and `locale` were missing from this list until this was flagged as a bug
205/// (BUGS.md C14/C15): two runs differing only in `--singleton-master-fraction`
206/// or `--locale` produced the exact same filename, so the second run silently
207/// overwrote the first's output.
208#[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    // BUGS.md C14/C15: every parameter that changes the output must be
229    // hashed in, or two runs differing only in this one collide on the same
230    // filename and silently overwrite each other.
231    only_entity.unwrap_or("").hash(&mut hasher);
232    format!("{}_{:x}", domain, hasher.finish())
233}
234
235/// Load and parse a domain schema JSON file.
236///
237/// On failure, the error message includes the path attempted and a hint listing
238/// available domains found in the same directory.
239pub fn load_schema(domain: &str, schemas_dir: &Path) -> Result<DomainSchema, String> {
240    let path = schemas_dir.join(format!("{domain}.json"));
241    // Case-sensitive exact match against the actual schema file names, so
242    // "KYC" is rejected the same way on every OS — on a case-insensitive
243    // filesystem (Windows), `read_to_string` alone would silently succeed
244    // for "KYC" via the "kyc.json" file, producing a different run hash
245    // than "kyc" for what the user intended to be the same domain.
246    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
263/// List available domain names (without .json extension) in a directory.
264fn 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/// Build a `PipelineConfig` from CLI / Python parameters and a parsed schema.
285///
286/// Validates `size >= 10`, distributes singleton/doublet/triplet records
287/// across entities, and assigns noise weights per difficulty setting.
288#[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    // `--only-entity`: the weight-based split below runs over this entity
348    // alone, so `total_unique`/`n_duplicates` land on it entirely instead of
349    // being shared out — no separate "recompute the ratio" step needed.
350    // Entities it FK-references get a small identifier-only pool appended
351    // after the main loop (see `pool_only` below); every other entity in
352    // the domain is simply absent from `entity_plans`.
353    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    // `--only-entity`: entities the target directly references via
463    // `fk_remaps` still need to exist enough to give it something plausible
464    // to point at, without paying for a full-size, fully-noised generation
465    // of an entity the run never writes. Generate just an identifier pool,
466    // capped at `FK_POOL_CAP` (same cap the FK-remap sampler itself already
467    // enforces, so a larger pool would never be consulted anyway).
468    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    // Hard negatives are per-entity (`HnSchema::entity_type`) — under
505    // `--only-entity` only the target's own hard-negative types apply; the
506    // rest belong to entities this run no longer writes.
507    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    // Scaled off `n_duplicates` (which itself scales with the tier's
514    // singleton/doublet fractions — light ~0.28*size, medium ~0.40*size,
515    // hell ~0.57*size) rather than off raw `size`. A flat `size`-based count
516    // makes `total_guaranteed_fp` roughly constant across tiers while
517    // `total_true_pairs` grows a lot with difficulty, which mechanically
518    // inflates `precision_max` for higher tiers regardless of how much
519    // noise they actually inject — enough to invert `light <= medium` on
520    // some schemas (e.g. aviation, where light/medium f1_max came out
521    // 0.828/0.833). The 0.125 constant keeps medium's hard-neg volume
522    // matching the old flat formula (0.40 * 0.125 = 0.05), so
523    // `hard_neg_ratio`'s documented "~1.5% of size at default 0.3" still
524    // holds at medium; light gets proportionally fewer, hell more.
525    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        // Windows filesystems are case-insensitive; load_schema must still
592        // reject "KYC" so run hashes stay consistent across OSes.
593        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        // Every entity must get at least the floor of 2 base records.
609        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    /// No-FK case (BUGS.md-style regression target, kyc's two entities
655    /// aren't FK-linked): `only_entity` should produce exactly one entity
656    /// plan, sized to the full `size`, with no pool-only entities appended.
657    #[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        // With only one active entity, all of `total_unique` (there's no
680        // weight split to share it with `legal_entity`) lands on it —
681        // comfortably more than the usual weight-shared slice would be.
682        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    /// FK case: aviation's `passenger` entity has no direct `fk_remaps` of
692    /// its own (verified against `schemas/aviation.json`), so `only_entity`
693    /// on it should still produce a single plan and no pool-only entities.
694    /// `flight`, which DOES reference `airline`/`aircraft`, is the case
695    /// that appends pool-only plans — covered by asserting the mechanism
696    /// generically below via `aircraft` (references `airline`).
697    #[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        // aircraft (target, written) + airline (pool-only, FK target).
716        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        // Pool-only entity is capped, not sized to the full run.
732        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    /// Regression: BUGS.md C14/C15 — `singleton_master_fraction` and
767    /// `locale` weren't hashed into the run ID, so two runs differing only
768    /// in one of these parameters (but producing genuinely different data)
769    /// got the exact same output filename, silently overwriting each other.
770    #[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    /// Regression guard for the `only_entity` hash input added alongside
781    /// `--only-entity`: two runs differing only in which entity is targeted
782    /// must not collide on the same run id (same class of bug as C14/C15).
783    #[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        // Unknown difficulty falls back to "medium".
816        assert_eq!(
817            default_singleton_master_fraction("bogus"),
818            default_singleton_master_fraction("medium")
819        );
820    }
821}