Skip to main content

dupehell_core/
schema.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6use std::collections::HashMap;
7use std::path::Path;
8
9use crate::pipeline::PipelineConfig;
10
11#[derive(serde::Deserialize)]
12pub struct DomainSchema {
13    pub entities: Vec<EntitySchema>,
14    pub hn_types: Vec<HnSchema>,
15}
16
17fn default_entity_weight() -> f64 {
18    1.0
19}
20
21/// Infer an entity's identifier column: prefer `{entity_name}_id`, then
22/// `id`, then the first column ending in `_id`.
23fn infer_identifier_col(entity: &EntitySchema) -> Option<String> {
24    let entity_id_name = format!("{}_id", entity.name);
25    let col_names: Vec<&str> = entity
26        .columns
27        .iter()
28        .filter_map(|c| c.get("name").and_then(|v| v.as_str()))
29        .collect();
30    if col_names.contains(&entity_id_name.as_str()) {
31        Some(entity_id_name)
32    } else if col_names.contains(&"id") {
33        Some("id".to_string())
34    } else {
35        col_names
36            .iter()
37            .find(|n| n.ends_with("_id"))
38            .map(|n| (*n).to_string())
39    }
40}
41
42#[derive(serde::Deserialize)]
43pub struct EntitySchema {
44    pub name: String,
45    pub columns: Vec<serde_json::Value>,
46    #[serde(default)]
47    pub fk_remaps: Vec<serde_json::Value>,
48    /// Relative population size of this entity within the domain, used to
49    /// split `total_unique` identities across entities (`build_pipeline_config`).
50    /// Defaults to `1.0` (every entity gets an equal share) when a schema
51    /// doesn't set it — every schema written before this field existed keeps
52    /// behaving exactly as before. Set explicit weights when entities aren't
53    /// realistically equal in population (e.g. aviation: far more
54    /// `passenger` identities than `airline` ones) — see `docs/` or ask
55    /// before picking numbers for a domain you don't know well; a wrong
56    /// weight is as misleading as the uniform default, just in the other
57    /// direction.
58    #[serde(default = "default_entity_weight")]
59    pub weight: f64,
60}
61
62#[derive(serde::Deserialize)]
63pub struct HnSchema {
64    pub entity_type: String,
65    pub config_json: String,
66}
67
68#[derive(Debug, Clone)]
69struct DifficultySettings {
70    singleton: f64,
71    doublet: f64,
72    /// Noise types active for this difficulty, weighted equally. Curated
73    /// explicitly per level (rather than derived from a truncated,
74    /// order-dependent countdown) so that "broad" noise types — `typo` and
75    /// `visual`, which both corrupt name+address+phone+company on the same
76    /// record at once (see `pipeline::noise_type_targets_column`) — never
77    /// end up a *larger* share of the mix at a lower difficulty than at a
78    /// higher one. A prior countdown-weight scheme caused exactly that:
79    /// medium (4 active types) put 70% of its noise budget on typo+visual,
80    /// vs. only ~42% for hell (8 active types), making medium duplicates
81    /// *more* likely to have every strong matching field wiped out at once
82    /// than hell duplicates, despite medium being meant as the easier tier.
83    noise_types: &'static [&'static str],
84    /// Number of *independent* noise passes applied to each duplicate copy
85    /// (see `pipeline::run_pipeline`'s dup-generation loop). Each additional
86    /// pass draws its own noise_type from `noise_types` and applies it on
87    /// top of the previous pass's result.
88    ///
89    /// This is the actual difficulty lever between tiers with different
90    /// `noise_types.len()` — not the category list itself. Each entry's
91    /// per-pass weight is `1 / noise_types.len()` (`build_pipeline_config`
92    /// below), so adding categories to make a tier "harder" *dilutes* every
93    /// existing category's weight, including whichever one happens to guard
94    /// a domain's single most reliable linkage column — a tier with more
95    /// categories active can paradoxically end up *easier* on some schemas
96    /// (measured: a "hard" tier here, folded into hell, kept coming out
97    /// easier than "hell" on several domains despite having fewer
98    /// categories; then a hand-picked category weighted twice fixed 2 of 6
99    /// domains tested but not the rest, since the dominant reliable column
100    /// isn't always the same category across schemas). Passes sidestep this
101    /// entirely: the probability a given column is touched at least once
102    /// across `P` independent passes is `1 - (1 - p)^P` where `p` is its
103    /// single-pass weight — strictly increasing in `P` regardless of how
104    /// small `p` is, so more passes can never make a tier easier the way
105    /// more categories can. `difficulty::estimate_difficulty` models this
106    /// exact same formula.
107    passes: usize,
108}
109
110const DIFFICULTY_MAP: [(&str, DifficultySettings); 3] = [
111    (
112        "light",
113        DifficultySettings {
114            singleton: 0.50,
115            doublet: 0.30,
116            noise_types: &["names", "dates"],
117            passes: 1,
118        },
119    ),
120    (
121        "medium",
122        DifficultySettings {
123            singleton: 0.30,
124            doublet: 0.40,
125            noise_types: &["typo", "names", "dates", "identifiers"],
126            // 3 passes: a column protected by exactly one of medium's 4
127            // equally-weighted categories (weight 1/4 -- e.g. a pure date
128            // column, only matched by "dates", not by "typo"'s broader
129            // predicate) is diluted well below light's worst case (2
130            // categories, weight 1/2). At passes=1 this let medium come out
131            // *easier* than light on 19/40 domains (found via
132            // scripts/validate_all_domains.py). 1 - (1 - 1/4)^3 ≈ 0.578 >
133            // light's single-pass 0.5, closing the gap for any category
134            // medium shares with light regardless of which column ends up
135            // being a given domain's escape hatch. Re-validated across all
136            // 40 domains with the same script after this change.
137            passes: 3,
138        },
139    ),
140    (
141        "hell",
142        DifficultySettings {
143            singleton: 0.10,
144            doublet: 0.20,
145            // `typo_aggressive` corrupts more of each string than plain
146            // `typo` (same target columns, more damage per hit — see
147            // `noise::typos`), and `unicode_pollution` gets its own bucket
148            // (rather than only ever showing up as one of `visual`'s five
149            // random sub-choices).
150            noise_types: &[
151                "typo_aggressive",
152                "visual",
153                "unicode_pollution",
154                "names",
155                "dates",
156                "identifiers",
157                "addresses",
158                "companies",
159                "extra",
160            ],
161            // 8 passes: medium was bumped from 1 to 3 passes to fix its own
162            // dilution against light (see the comment on medium's `passes`
163            // above), which raised medium's worst-case single-category
164            // touch probability to 1-(1-1/4)^3 ≈ 0.578. With 9
165            // equally-weighted categories (weight 1/9 each), hell needs
166            // 1-(1-1/9)^passes >= 0.578 to stay at least as hard as medium
167            // on any column protected by only one category -- solving gives
168            // passes >= 7.33, rounded up to 8 (1-(1-1/9)^8 ≈ 0.617).
169            // Re-validated across all 40 domains via
170            // scripts/validate_all_domains.py (0 crashes, 0 monotonicity
171            // violations) after this change.
172            passes: 8,
173        },
174    ),
175];
176
177pub fn default_singleton_master_fraction(difficulty: &str) -> f64 {
178    difficulty_settings(difficulty).singleton
179}
180
181fn difficulty_settings(difficulty: &str) -> DifficultySettings {
182    DIFFICULTY_MAP
183        .iter()
184        .find(|(name, _)| *name == difficulty)
185        .map(|(_, s)| s.clone())
186        .unwrap_or_else(|| DIFFICULTY_MAP[1].1.clone())
187}
188
189/// Generate a domain-unique run ID based on the current Unix timestamp (hex).
190pub fn chrono_now() -> String {
191    let start = std::time::SystemTime::now()
192        .duration_since(std::time::UNIX_EPOCH)
193        .unwrap_or_default();
194    format!("{:x}", start.as_secs())
195}
196
197/// Deterministic run ID derived from generation parameters, so the same
198/// (domain, size, seed, difficulty, hard_neg_ratio, singleton_master_fraction,
199/// locale) always produces the same output filename regardless of output
200/// format (IPC vs Parquet) or how many times it's run — and, just as
201/// important, a *different* filename whenever any of these parameters
202/// differ, since they all affect the generated data. `singleton_master_fraction`
203/// and `locale` were missing from this list until this was flagged as a bug
204/// (BUGS.md C14/C15): two runs differing only in `--singleton-master-fraction`
205/// or `--locale` produced the exact same filename, so the second run silently
206/// overwrote the first's output.
207#[allow(clippy::too_many_arguments)]
208pub fn deterministic_run_id(
209    domain: &str,
210    size: usize,
211    seed: u64,
212    difficulty: &str,
213    hard_neg_ratio: f64,
214    singleton_master_fraction: f64,
215    locale: &str,
216    only_entity: Option<&str>,
217    chunk_size: Option<usize>,
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    // `--chunk-size` doesn't change the final dataset's *content* (chunked
233    // and non-chunked runs of the same total size are meant to be
234    // equivalent data, modulo per-chunk RNG independence) but it DOES
235    // change the sequence of per-chunk seeds actually used internally
236    // (`seed + chunk_idx` per chunk vs a single unbroken `seed`), so two
237    // runs with the same total size/seed but a different chunking must
238    // never collide on the same output filename.
239    chunk_size.unwrap_or(0).hash(&mut hasher);
240    format!("{}_{:x}", domain, hasher.finish())
241}
242
243/// Load and parse a domain schema JSON file.
244///
245/// On failure, the error message includes the path attempted and a hint listing
246/// available domains found in the same directory.
247pub fn load_schema(domain: &str, schemas_dir: &Path) -> Result<DomainSchema, String> {
248    let path = schemas_dir.join(format!("{domain}.json"));
249    // Case-sensitive exact match against the actual schema file names, so
250    // "KYC" is rejected the same way on every OS — on a case-insensitive
251    // filesystem (Windows), `read_to_string` alone would silently succeed
252    // for "KYC" via the "kyc.json" file, producing a different run hash
253    // than "kyc" for what the user intended to be the same domain.
254    let available = available_domain_names(schemas_dir);
255    if !available.iter().any(|d| d == domain) {
256        let hint = if available.is_empty() {
257            "no schemas found".to_string()
258        } else {
259            available.join(", ")
260        };
261        return Err(format!(
262            "schema file not found for domain '{domain}' at {path:?}. \
263             Available domains ({hint})"
264        ));
265    }
266    let data =
267        std::fs::read_to_string(&path).map_err(|e| format!("cannot read schema {path:?}: {e}"))?;
268    serde_json::from_str(&data).map_err(|e| format!("cannot parse schema {domain}.json: {e}"))
269}
270
271/// List available domain names (without .json extension) in a directory.
272fn available_domain_names(dir: &Path) -> Vec<String> {
273    match std::fs::read_dir(dir) {
274        Ok(entries) => entries
275            .filter_map(|e| e.ok())
276            .filter(|e| {
277                e.path()
278                    .extension()
279                    .map(|ext| ext == "json")
280                    .unwrap_or(false)
281            })
282            .filter_map(|e| {
283                e.path()
284                    .file_stem()
285                    .map(|s| s.to_string_lossy().to_string())
286            })
287            .collect(),
288        Err(_) => Vec::new(),
289    }
290}
291
292/// Build a `PipelineConfig` from CLI / Python parameters and a parsed schema.
293///
294/// Validates `size >= 10`, distributes singleton/doublet/triplet records
295/// across entities, and assigns noise weights per difficulty setting.
296#[allow(clippy::too_many_arguments)]
297pub fn build_pipeline_config(
298    domain: &str,
299    size: usize,
300    seed: u64,
301    difficulty: &str,
302    hard_neg_ratio: f64,
303    singleton_master_fraction: f64,
304    schema: &DomainSchema,
305    run_id: &str,
306    output_format: &str,
307    graph_enabled: bool,
308    graph_format: &str,
309    only_entity: Option<&str>,
310) -> Result<PipelineConfig, String> {
311    if size < 10 {
312        return Err(format!("size must be >= 10, got {size}"));
313    }
314    if schema.entities.is_empty() {
315        return Err(format!("schema for domain '{domain}' has no entities"));
316    }
317    if !(0.0..=1.0).contains(&singleton_master_fraction) {
318        return Err(format!(
319            "singleton_master_fraction must be in [0.0, 1.0], got {singleton_master_fraction}"
320        ));
321    }
322    if let Some(target) = only_entity
323        && !schema.entities.iter().any(|e| e.name == target)
324    {
325        let available: Vec<&str> = schema.entities.iter().map(|e| e.name.as_str()).collect();
326        return Err(format!(
327            "only_entity '{target}' not found in domain '{domain}'; available entities: {}",
328            available.join(", ")
329        ));
330    }
331    let ds = difficulty_settings(difficulty);
332    let total = size;
333
334    let n_singleton = (total as f64 * singleton_master_fraction) as usize;
335    let n_doublet_float = total as f64 * ds.doublet;
336    let mut n_doublet = n_doublet_float as usize;
337    if !n_doublet.is_multiple_of(2) {
338        n_doublet -= 1;
339    }
340    if n_singleton + n_doublet > total {
341        return Err(format!(
342            "singleton_master_fraction {singleton_master_fraction} leaves no room for this \
343             difficulty's doublet share ({:.2}); reduce singleton_master_fraction",
344            ds.doublet
345        ));
346    }
347    let mut n_triplet = total - n_singleton - n_doublet;
348    let r = n_triplet % 3;
349    if r != 0 {
350        n_triplet -= r;
351    }
352    let total_unique = n_singleton + n_doublet / 2 + n_triplet / 3;
353    let n_duplicates = total.max(total_unique) - total_unique;
354
355    // `--only-entity`: the weight-based split below runs over this entity
356    // alone, so `total_unique`/`n_duplicates` land on it entirely instead of
357    // being shared out — no separate "recompute the ratio" step needed.
358    // Entities it FK-references get a small identifier-only pool appended
359    // after the main loop (see `pool_only` below); every other entity in
360    // the domain is simply absent from `entity_plans`.
361    let active_entities: Vec<&EntitySchema> = match only_entity {
362        Some(target) => schema
363            .entities
364            .iter()
365            .filter(|e| e.name == target)
366            .collect(),
367        None => schema.entities.iter().collect(),
368    };
369
370    let total_ratio: f64 = active_entities.iter().map(|e| e.weight).sum::<f64>();
371    let raw_floats: Vec<(&str, f64)> = active_entities
372        .iter()
373        .map(|e| {
374            (
375                e.name.as_str(),
376                total_unique as f64 * e.weight / total_ratio,
377            )
378        })
379        .collect();
380
381    let mut floor_map: HashMap<&str, usize> = HashMap::new();
382    for (name, r) in &raw_floats {
383        floor_map.insert(name, r.max(2.0) as usize);
384    }
385    let floor_sum: usize = floor_map.values().sum();
386    let need = total_unique.max(floor_sum) - floor_sum;
387    if need > 0 {
388        let mut remainders: Vec<(&str, f64)> = raw_floats
389            .iter()
390            .map(|(n, r)| (*n, r - r.floor()))
391            .collect();
392        remainders.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
393        for (name, _) in remainders.iter().take(need) {
394            *floor_map.get_mut(name).unwrap_or(&mut 0) += 1;
395        }
396    }
397
398    let dup_ratios: Vec<(&str, f64)> = active_entities
399        .iter()
400        .map(|e| {
401            (
402                e.name.as_str(),
403                *floor_map.get(e.name.as_str()).unwrap_or(&2) as f64 / total_unique as f64,
404            )
405        })
406        .collect();
407    let mut dup_floor: HashMap<&str, usize> = HashMap::new();
408    for (name, r) in &dup_ratios {
409        dup_floor.insert(name, (n_duplicates as f64 * r) as usize);
410    }
411    let dup_sum: usize = dup_floor.values().sum();
412    let dup_need = n_duplicates.max(dup_sum) - dup_sum;
413    if dup_need > 0 {
414        let mut remainders: Vec<(&str, f64)> = dup_ratios
415            .iter()
416            .map(|(n, r)| (*n, r - r.floor()))
417            .collect();
418        remainders.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
419        for (name, _) in remainders.iter().take(dup_need) {
420            *dup_floor.get_mut(name).unwrap_or(&mut 0) += 1;
421        }
422    }
423
424    let noise_count = ds.noise_types.len();
425    let noise_weights: Vec<f64> = vec![1.0 / noise_count as f64; noise_count];
426
427    let mut entity_plans = Vec::new();
428    for entity in &active_entities {
429        let n_base = *floor_map.get(entity.name.as_str()).unwrap_or(&2);
430        let n_dup = *dup_floor.get(entity.name.as_str()).unwrap_or(&0);
431
432        let mut noise_entries = Vec::new();
433        if n_dup > 0 {
434            let mut counts: Vec<usize> = noise_weights
435                .iter()
436                .map(|w| (w * n_dup as f64) as usize)
437                .collect();
438            let count_sum: usize = counts.iter().sum();
439            if count_sum < n_dup {
440                *counts.last_mut().unwrap_or(&mut 0) += n_dup - count_sum;
441            }
442            for (i, noise_type) in ds.noise_types.iter().enumerate() {
443                if counts[i] == 0 {
444                    continue;
445                }
446                noise_entries.push(serde_json::json!({
447                    "noise_type": noise_type,
448                    "columns": [],
449                    "count": counts[i],
450                }));
451            }
452        }
453
454        let identifier_col = infer_identifier_col(entity);
455
456        let columns_json = serde_json::to_string(&entity.columns)
457            .map_err(|e| format!("serialize columns: {e}"))?;
458
459        entity_plans.push(serde_json::json!({
460            "name": entity.name,
461            "n_base": n_base,
462            "n_dup": n_dup,
463            "identifier_col": identifier_col,
464            "columns_json": columns_json,
465            "noise_types": noise_entries,
466            "fk_remaps": entity.fk_remaps,
467        }));
468    }
469
470    // `--only-entity`: entities the target directly references via
471    // `fk_remaps` still need to exist enough to give it something plausible
472    // to point at, without paying for a full-size, fully-noised generation
473    // of an entity the run never writes. Generate just an identifier pool,
474    // capped at `FK_POOL_CAP` (same cap the FK-remap sampler itself already
475    // enforces, so a larger pool would never be consulted anyway).
476    if let Some(target) = only_entity {
477        let target_schema = schema
478            .entities
479            .iter()
480            .find(|e| e.name == target)
481            .expect("only_entity validated to exist above");
482        let mut pool_target_names: Vec<&str> = target_schema
483            .fk_remaps
484            .iter()
485            .filter_map(|r| r.get("target_entity").and_then(|v| v.as_str()))
486            .filter(|&t| t != target)
487            .collect();
488        pool_target_names.sort_unstable();
489        pool_target_names.dedup();
490
491        let pool_n = crate::pipeline::FK_POOL_CAP.min(total_unique.max(2));
492        for name in pool_target_names {
493            let Some(entity) = schema.entities.iter().find(|e| e.name == name) else {
494                continue;
495            };
496            let identifier_col = infer_identifier_col(entity);
497            let columns_json = serde_json::to_string(&entity.columns)
498                .map_err(|e| format!("serialize columns: {e}"))?;
499            entity_plans.push(serde_json::json!({
500                "name": entity.name,
501                "n_base": pool_n,
502                "n_dup": 0,
503                "identifier_col": identifier_col,
504                "columns_json": columns_json,
505                "noise_types": [],
506                "fk_remaps": entity.fk_remaps,
507                "pool_only": true,
508            }));
509        }
510    }
511
512    // Hard negatives are per-entity (`HnSchema::entity_type`) — under
513    // `--only-entity` only the target's own hard-negative types apply; the
514    // rest belong to entities this run no longer writes.
515    let active_hn_types: Vec<&HnSchema> = schema
516        .hn_types
517        .iter()
518        .filter(|hn| only_entity.is_none_or(|t| hn.entity_type == t))
519        .collect();
520
521    // Scaled off `n_duplicates` (which itself scales with the tier's
522    // singleton/doublet fractions — light ~0.28*size, medium ~0.40*size,
523    // hell ~0.57*size) rather than off raw `size`. A flat `size`-based count
524    // makes `total_guaranteed_fp` roughly constant across tiers while
525    // `total_true_pairs` grows a lot with difficulty, which mechanically
526    // inflates `precision_max` for higher tiers regardless of how much
527    // noise they actually inject — enough to invert `light <= medium` on
528    // some schemas (e.g. aviation, where light/medium f1_max came out
529    // 0.828/0.833). The 0.125 constant keeps medium's hard-neg volume
530    // matching the old flat formula (0.40 * 0.125 = 0.05), so
531    // `hard_neg_ratio`'s documented "~1.5% of size at default 0.3" still
532    // holds at medium; light gets proportionally fewer, hell more.
533    let n_hard_neg = (n_duplicates as f64 * hard_neg_ratio * 0.125) as usize;
534    let hn_per_type = n_hard_neg / active_hn_types.len().max(1);
535    let hard_neg_types: Vec<serde_json::Value> = active_hn_types
536        .iter()
537        .map(|hn| {
538            serde_json::json!({
539                "entity_type": hn.entity_type,
540                "config_json": hn.config_json,
541                "count": hn_per_type,
542            })
543        })
544        .collect();
545
546    let config = serde_json::json!({
547        "domain": domain,
548        "size": size,
549        "seed": seed,
550        "difficulty": difficulty,
551        "output_format": output_format,
552        "run_id": run_id,
553        "entity_plans": entity_plans,
554        "hard_neg_types": hard_neg_types,
555        "hard_neg_ratio": hard_neg_ratio,
556        "noise_passes": ds.passes,
557        "graph_enabled": graph_enabled,
558        "graph_format": graph_format,
559    });
560
561    serde_json::from_value(config).map_err(|e| format!("build PipelineConfig: {e}"))
562}
563
564#[cfg(test)]
565mod tests {
566    use super::*;
567
568    fn schemas_dir() -> std::path::PathBuf {
569        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("schemas")
570    }
571
572    fn kyc_schema() -> DomainSchema {
573        load_schema("kyc", &schemas_dir()).expect("load kyc.json")
574    }
575
576    fn aviation_schema() -> DomainSchema {
577        load_schema("aviation", &schemas_dir()).expect("load aviation.json")
578    }
579
580    #[test]
581    fn test_load_schema_known_domain() {
582        let schema = kyc_schema();
583        assert_eq!(schema.entities.len(), 2);
584        assert!(schema.entities.iter().any(|e| e.name == "natural_person"));
585        assert!(!schema.hn_types.is_empty());
586    }
587
588    #[test]
589    fn test_load_schema_unknown_domain_lists_available() {
590        let err = load_schema("not-a-real-domain", &schemas_dir())
591            .err()
592            .unwrap();
593        assert!(err.contains("not-a-real-domain"));
594        assert!(err.contains("kyc"));
595    }
596
597    #[test]
598    fn test_load_schema_case_sensitive() {
599        // Windows filesystems are case-insensitive; load_schema must still
600        // reject "KYC" so run hashes stay consistent across OSes.
601        let err = load_schema("KYC", &schemas_dir()).err().unwrap();
602        assert!(err.contains("KYC"));
603    }
604
605    #[test]
606    fn test_build_pipeline_config_basic() {
607        let schema = kyc_schema();
608        let config = build_pipeline_config(
609            "kyc", 1000, 42, "medium", 0.1, 0.3, &schema, "kyc_test", "parquet", false, "parquet",
610            None,
611        )
612        .expect("build config");
613        assert_eq!(config.domain, "kyc");
614        assert_eq!(config.size, 1000);
615        assert_eq!(config.entity_plans.len(), 2);
616        // Every entity must get at least the floor of 2 base records.
617        assert!(config.entity_plans.iter().all(|p| p.n_base >= 2));
618        assert!(!config.hard_neg_types.is_empty());
619    }
620
621    #[test]
622    fn test_build_pipeline_config_rejects_size_below_10() {
623        let schema = kyc_schema();
624        let err = build_pipeline_config(
625            "kyc", 5, 42, "medium", 0.1, 0.3, &schema, "kyc_test", "parquet", false, "parquet",
626            None,
627        )
628        .unwrap_err();
629        assert!(err.contains("size must be >= 10"));
630    }
631
632    #[test]
633    fn test_build_pipeline_config_rejects_invalid_singleton_fraction() {
634        let schema = kyc_schema();
635        let err = build_pipeline_config(
636            "kyc", 1000, 42, "medium", 0.1, 1.5, &schema, "kyc_test", "parquet", false, "parquet",
637            None,
638        )
639        .unwrap_err();
640        assert!(err.contains("singleton_master_fraction"));
641    }
642
643    #[test]
644    fn test_build_pipeline_config_deterministic() {
645        let schema = kyc_schema();
646        let build = || {
647            build_pipeline_config(
648                "kyc", 1000, 42, "hell", 0.1, 0.2, &schema, "kyc_test", "parquet", false,
649                "parquet", None,
650            )
651            .expect("build config")
652        };
653        let a = build();
654        let b = build();
655        assert_eq!(a.entity_plans.len(), b.entity_plans.len());
656        for (pa, pb) in a.entity_plans.iter().zip(b.entity_plans.iter()) {
657            assert_eq!(pa.n_base, pb.n_base);
658            assert_eq!(pa.noise_types.len(), pb.noise_types.len());
659        }
660    }
661
662    /// No-FK case (BUGS.md-style regression target, kyc's two entities
663    /// aren't FK-linked): `only_entity` should produce exactly one entity
664    /// plan, sized to the full `size`, with no pool-only entities appended.
665    #[test]
666    fn test_build_pipeline_config_only_entity_no_fk() {
667        let schema = kyc_schema();
668        let config = build_pipeline_config(
669            "kyc",
670            1000,
671            42,
672            "medium",
673            0.1,
674            0.3,
675            &schema,
676            "kyc_test",
677            "parquet",
678            false,
679            "parquet",
680            Some("natural_person"),
681        )
682        .expect("build config");
683        assert_eq!(config.entity_plans.len(), 1);
684        let plan = &config.entity_plans[0];
685        assert_eq!(plan.name, "natural_person");
686        assert!(!plan.pool_only);
687        // With only one active entity, all of `total_unique` (there's no
688        // weight split to share it with `legal_entity`) lands on it —
689        // comfortably more than the usual weight-shared slice would be.
690        assert!(plan.n_base > 100);
691        assert!(
692            config
693                .hard_neg_types
694                .iter()
695                .all(|hn| hn.entity_type == "natural_person")
696        );
697    }
698
699    /// FK case: aviation's `passenger` entity has no direct `fk_remaps` of
700    /// its own (verified against `schemas/aviation.json`), so `only_entity`
701    /// on it should still produce a single plan and no pool-only entities.
702    /// `flight`, which DOES reference `airline`/`aircraft`, is the case
703    /// that appends pool-only plans — covered by asserting the mechanism
704    /// generically below via `aircraft` (references `airline`).
705    #[test]
706    fn test_build_pipeline_config_only_entity_with_fk_pool() {
707        let schema = aviation_schema();
708        let config = build_pipeline_config(
709            "aviation",
710            10_000,
711            42,
712            "medium",
713            0.1,
714            0.3,
715            &schema,
716            "aviation_test",
717            "parquet",
718            false,
719            "parquet",
720            Some("aircraft"),
721        )
722        .expect("build config");
723        // aircraft (target, written) + airline (pool-only, FK target).
724        assert_eq!(config.entity_plans.len(), 2);
725        let target = config
726            .entity_plans
727            .iter()
728            .find(|p| p.name == "aircraft")
729            .expect("aircraft plan present");
730        assert!(!target.pool_only);
731        assert!(target.n_base > 1_000);
732        let pool = config
733            .entity_plans
734            .iter()
735            .find(|p| p.name == "airline")
736            .expect("airline pool-only plan present");
737        assert!(pool.pool_only);
738        assert!(pool.noise_types.is_empty());
739        // Pool-only entity is capped, not sized to the full run.
740        assert!(pool.n_base <= crate::pipeline::FK_POOL_CAP);
741    }
742
743    #[test]
744    fn test_build_pipeline_config_only_entity_unknown_name() {
745        let schema = kyc_schema();
746        let err = build_pipeline_config(
747            "kyc",
748            1000,
749            42,
750            "medium",
751            0.1,
752            0.3,
753            &schema,
754            "kyc_test",
755            "parquet",
756            false,
757            "parquet",
758            Some("not_a_real_entity"),
759        )
760        .unwrap_err();
761        assert!(err.contains("not_a_real_entity"));
762        assert!(err.contains("natural_person"));
763    }
764
765    #[test]
766    fn test_deterministic_run_id_stable_and_sensitive() {
767        let a = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
768        let b = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
769        assert_eq!(a, b);
770        let c = deterministic_run_id("kyc", 1000, 43, "medium", 0.1, 0.3, "en", None, None);
771        assert_ne!(a, c);
772    }
773
774    /// Regression: BUGS.md C14/C15 — `singleton_master_fraction` and
775    /// `locale` weren't hashed into the run ID, so two runs differing only
776    /// in one of these parameters (but producing genuinely different data)
777    /// got the exact same output filename, silently overwriting each other.
778    #[test]
779    fn test_deterministic_run_id_sensitive_to_singleton_fraction_and_locale() {
780        let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
781        let diff_fraction =
782            deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.5, "en", None, None);
783        let diff_locale =
784            deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "fr", None, None);
785        assert_ne!(base, diff_fraction);
786        assert_ne!(base, diff_locale);
787        assert_ne!(diff_fraction, diff_locale);
788    }
789
790    /// Regression guard for the `only_entity` hash input added alongside
791    /// `--only-entity`: two runs differing only in which entity is targeted
792    /// must not collide on the same run id (same class of bug as C14/C15).
793    #[test]
794    fn test_deterministic_run_id_sensitive_to_only_entity() {
795        let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
796        let only_a = deterministic_run_id(
797            "kyc",
798            1000,
799            42,
800            "medium",
801            0.1,
802            0.3,
803            "en",
804            Some("natural_person"),
805            None,
806        );
807        let only_b = deterministic_run_id(
808            "kyc",
809            1000,
810            42,
811            "medium",
812            0.1,
813            0.3,
814            "en",
815            Some("legal_entity"),
816            None,
817        );
818        assert_ne!(base, only_a);
819        assert_ne!(base, only_b);
820        assert_ne!(only_a, only_b);
821    }
822
823    /// Regression guard for the `chunk_size` hash input added alongside
824    /// `--chunk-size`: two runs with the same total size but a different
825    /// chunking must not collide on the same run id, since the per-chunk
826    /// seed sequence actually used internally differs (same class of bug
827    /// as C14/C15).
828    #[test]
829    fn test_deterministic_run_id_sensitive_to_chunk_size() {
830        let base = deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, None);
831        let chunked_a =
832            deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, Some(500));
833        let chunked_b =
834            deterministic_run_id("kyc", 1000, 42, "medium", 0.1, 0.3, "en", None, Some(250));
835        assert_ne!(base, chunked_a);
836        assert_ne!(base, chunked_b);
837        assert_ne!(chunked_a, chunked_b);
838    }
839
840    #[test]
841    fn test_default_singleton_master_fraction_known_and_unknown() {
842        assert_eq!(default_singleton_master_fraction("light"), 0.50);
843        assert_eq!(default_singleton_master_fraction("hell"), 0.10);
844        // Unknown difficulty falls back to "medium".
845        assert_eq!(
846            default_singleton_master_fraction("bogus"),
847            default_singleton_master_fraction("medium")
848        );
849    }
850}