Skip to main content

dupehell_core/
pipeline.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::sync::Arc;
9
10use arrow::array::{
11    Array, ArrayBuilder, ArrayRef, AsArray, StringArray, UInt64Array, UInt64Builder,
12};
13use arrow::datatypes::{DataType, Field, Schema};
14use arrow::record_batch::RecordBatch;
15use rayon::prelude::*;
16use rustc_hash::FxHashMap;
17use serde::Deserialize;
18
19use crate::context::Context;
20use crate::rng::Rng;
21
22// ── Config types ───────────────────────────────────────────────────────────
23
24#[derive(Debug, Deserialize)]
25pub struct PipelineConfig {
26    pub domain: String,
27    pub size: usize,
28    pub seed: u64,
29    pub difficulty: String,
30    pub output_format: String,
31    pub run_id: String,
32
33    pub entity_plans: Vec<EntityPlan>,
34    pub hard_neg_types: Vec<HnTypeConfig>,
35
36    pub hard_neg_ratio: f64,
37
38    /// Number of independent noise passes applied to each duplicate copy —
39    /// see the doc comment on `schema::DifficultySettings::passes` for why
40    /// this (not the noise_types category count) is the actual difficulty
41    /// lever between tiers.
42    pub noise_passes: usize,
43
44    pub graph_enabled: bool,
45    pub graph_format: String,
46}
47
48#[derive(Debug, Deserialize)]
49pub struct EntityPlan {
50    pub name: String,
51    pub n_base: usize,
52    pub identifier_col: Option<String>,
53    pub columns_json: String,
54    pub noise_types: Vec<NoisePlanEntry>,
55    pub fk_remaps: Vec<FkRemapEntry>,
56    /// Set by `schema::build_pipeline_config` for the FK-target entities
57    /// added when `--only-entity` is active (see `FK_POOL_CAP`): this
58    /// entity's rows are generated only far enough to build its identifier
59    /// pool for other entities' FK remaps, never noised, written to the
60    /// dataset/graph output, or fed to ground truth. Always `false` on the
61    /// normal (no `--only-entity`) path.
62    #[serde(default)]
63    pub pool_only: bool,
64}
65
66/// Cap on how many identifiers of an entity are retained as a pool for other
67/// entities' FK remaps to sample from (`fk_remap.rs` samples uniformly via
68/// `take`, so it never needs to see every identifier the entity produced).
69/// Also used by `schema::build_pipeline_config` to size the pool-only
70/// entities it synthesizes for `--only-entity` (no need to generate more
71/// rows of an FK target than this pool will ever draw from).
72pub(crate) const FK_POOL_CAP: usize = 200_000;
73
74#[derive(Debug, Deserialize)]
75pub struct NoisePlanEntry {
76    pub noise_type: String,
77    pub columns: Vec<String>,
78    pub count: usize,
79}
80
81#[derive(Debug, Deserialize)]
82pub struct FkRemapEntry {
83    pub source_col: String,
84    pub target_entity: String,
85    /// Weight written to the `weight` column of `fk`-typed graph edges
86    /// produced from this relation (`--graph`). Defaults to 1.0 for schemas
87    /// that don't declare it, preserving prior behavior.
88    #[serde(default = "default_fk_weight")]
89    pub fk_weight: f64,
90}
91
92fn default_fk_weight() -> f64 {
93    1.0
94}
95
96#[derive(Debug, Deserialize)]
97pub struct HnTypeConfig {
98    pub entity_type: String,
99    pub config_json: String,
100    pub count: usize,
101}
102
103#[derive(Debug)]
104pub struct PipelineOutput {
105    pub output_files: Vec<String>,
106    pub gt_file: String,
107    pub stats: PipelineStats,
108    pub nodes: Option<String>,
109    pub edges: Option<String>,
110}
111
112#[derive(Debug)]
113pub struct PipelineStats {
114    pub total_records: usize,
115    /// Cluster members that are genuinely byte-for-byte identical to their
116    /// master. See `PipelineStats::fuzzy_dups` for the noised counterpart.
117    pub exact_dups: usize,
118    /// Cluster members that are duplicates of a master but differ from it on
119    /// at least one column (the assigned noise actually changed something).
120    pub fuzzy_dups: usize,
121    pub hard_negs: usize,
122    pub uniques: usize,
123    pub masters: usize,
124}
125
126// ── Fixed-width ASCII IDs ────────────────────────────────────────────────
127
128// `record_id`/`pad` are pure functions of the row index — "R-" + 13 digits
129// and 13 digits respectively. Earlier this was a struct precomputing both
130// for every id up front (bounded per-2-big-allocations, better than one
131// `String` per id, but still O(total_ids) resident for the whole run — at
132// 200M+ ids that's several GB held for nothing). Computing on demand, only
133// for the batch currently being written, removes that fixed cost entirely;
134// the transient `Vec<String>` built per batch is bounded by `BATCH_SIZE`
135// and freed right after use, same as `batch_mids`/`dup_mids_buf`/`hn_mids`
136// elsewhere in this file.
137const RID_LEN: usize = 15; // "R-" + 13 digits
138const PAD_LEN: usize = 13; // 13 digits
139
140#[inline]
141pub(crate) fn record_id_string(i: usize) -> String {
142    let mut buf = [0u8; RID_LEN];
143    buf[0] = b'R';
144    buf[1] = b'-';
145    let mut n = i;
146    for j in (2..RID_LEN).rev() {
147        buf[j] = b'0' + (n % 10) as u8;
148        n /= 10;
149    }
150    String::from_utf8(buf.to_vec()).expect("record_id_string: buffer is 'R', '-' and ASCII digits")
151}
152
153#[inline]
154pub(crate) fn pad_string(i: usize) -> String {
155    let mut buf = [0u8; PAD_LEN];
156    let mut n = i;
157    for j in (0..PAD_LEN).rev() {
158        buf[j] = b'0' + (n % 10) as u8;
159        n /= 10;
160    }
161    String::from_utf8(buf.to_vec()).expect("pad_string: buffer is always ASCII digits")
162}
163
164pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
165    range.map(record_id_string).collect()
166}
167
168/// Inverse of `record_id_string`: parses `"R-0000000000007"` back to `7`.
169/// Returns `None` if `rid` isn't exactly `RID_LEN` bytes starting with
170/// `"R-"` followed by `PAD_LEN` ASCII digits (never expected in practice —
171/// every `record_id` in a duplicated cluster is produced by
172/// `record_id_string`/`record_id_strs`, never hand-built).
173pub(crate) fn parse_record_idx(rid: &str) -> Option<u64> {
174    if rid.len() != RID_LEN || !rid.starts_with("R-") {
175        return None;
176    }
177    rid[2..].parse().ok()
178}
179
180/// Packs a `master_id` of the fixed shape `"{entity_prefix}-{pad_string}"`
181/// (e.g. `"E00042-0000000001234"`) into a single `u64`: `entity_idx` in the
182/// high bits, `local_idx` in the low 44 bits. `entity_idx` is at most 5
183/// digits (`entity_prefix`'s format), so it fits comfortably above bit 44;
184/// `local_idx` is at most `PAD_LEN` (13) digits, which fits in 44 bits
185/// (2^44 ≈ 1.76e13). Preserves the same ordering as a lexicographic
186/// comparison of the original fixed-width string, since both segments are
187/// zero-padded to a constant width — callers that need the old
188/// `master_ids.sort()` determinism can just sort/compare the packed `u64`.
189/// Returns `None` if `mid` isn't `"E" + 5 digits + "-" + PAD_LEN digits`
190/// (never expected for a genuine duplicated master_id — see
191/// `gt::GtAccumulator::push_dup_batch`, the only source of `cluster_map`
192/// keys, which always writes this exact shape).
193pub(crate) fn pack_master_key(mid: &str) -> Option<u64> {
194    let bytes = mid.as_bytes();
195    if bytes.len() != 1 + 5 + 1 + PAD_LEN || bytes[0] != b'E' || bytes[6] != b'-' {
196        return None;
197    }
198    let entity_idx: u64 = mid[1..6].parse().ok()?;
199    let local_idx: u64 = mid[7..7 + PAD_LEN].parse().ok()?;
200    Some((entity_idx << 44) | local_idx)
201}
202
203/// Splits `total` proportionally across `weights` (largest-remainder
204/// method), so the sum of the result is always exactly `total`.
205///
206/// Used to spread each noise_type's planned duplicate count across every
207/// batch of masters in proportion to batch size, instead of assigning it
208/// all to a single batch — see the doc comment on the "Dups" section in
209/// `run_pipeline_with_progress` for why that matters.
210fn distribute_by_weight(total: usize, weights: &[usize]) -> Vec<usize> {
211    let weight_sum: usize = weights.iter().sum();
212    if weight_sum == 0 || total == 0 {
213        return vec![0; weights.len()];
214    }
215    let raw: Vec<f64> = weights
216        .iter()
217        .map(|&w| total as f64 * w as f64 / weight_sum as f64)
218        .collect();
219    let mut counts: Vec<usize> = raw.iter().map(|&r| r as usize).collect();
220    let assigned: usize = counts.iter().sum();
221    let mut remainder = total - assigned;
222    let mut order: Vec<usize> = (0..counts.len()).collect();
223    order.sort_by(|&a, &b| {
224        let fa = raw[a] - counts[a] as f64;
225        let fb = raw[b] - counts[b] as f64;
226        fb.partial_cmp(&fa).unwrap_or(std::cmp::Ordering::Equal)
227    });
228    for &i in &order {
229        if remainder == 0 {
230            break;
231        }
232        counts[i] += 1;
233        remainder -= 1;
234    }
235    counts
236}
237
238// ── Entity prefix ──────────────────────────────────────────────────────────
239
240/// Per-entity master_id prefix, derived from the entity's position in
241/// `config.entity_plans` rather than a hash of its name: a name hash mod
242/// 100000 had a real (if small) chance of two entities in the same domain
243/// colliding on the same prefix, which would make `gt.rs`'s master_id
244/// counting conflate masters from different entity types. An index is
245/// unique by construction — zero collision risk regardless of domain size.
246pub(crate) fn entity_prefix(entity_idx: usize) -> String {
247    format!("E{:05}", entity_idx)
248}
249
250// ── FK / HN pool types ─────────────────────────────────────────────────────
251
252type FkPoolMap = HashMap<String, RecordBatch>;
253
254struct HnPool {
255    batch: RecordBatch,
256    total_count: usize,
257    // Global record_ids aligned positionally with `batch` rows (only populated
258    // when graph output is enabled). Used to emit `hard_neg` edges.
259    record_ids: Vec<String>,
260}
261
262// ── Noise column matching ─────────────────────────────────────────────────
263
264/// Column-name fragments for columns holding a person's name as free text,
265/// even when the column isn't literally called `*_name` (e.g. `operator`,
266/// `technician`, populated from the `first_name` pool across several
267/// domain schemas — see the 40-domain schema audit in the 2026-07 session).
268const PERSON_NAME_WORDS: &[&str] = &[
269    "name",
270    // Not bare "first"/"last": those also matched non-name columns like
271    // `first_seen`/`last_active`/`last_login`/`last_updated`, sending them
272    // through char-level typo/homoglyph noise and injecting letters into
273    // what are actually timestamps (BUGS.md C13). Every legitimate name
274    // column across the 40 domain schemas is literally `*first_name*` /
275    // `*last_name*` (e.g. `agent_first_name`, `pat_first_name`), so this
276    // loses no real coverage.
277    "first_name",
278    "last_name",
279    "given",
280    "family",
281    "middle",
282    "operator",
283    "technician",
284    "inspector",
285    "claimant",
286    "consignee",
287    "investigator",
288    "collected_by",
289    "performed_by",
290    "assigned_to",
291];
292
293/// Column-name fragments for columns holding a company/organization name,
294/// even when the column isn't literally called `*_name`/`company*` (e.g.
295/// `supplier`, `manufacturer`, populated from the `company` pool across
296/// several domain schemas — see the 40-domain schema audit in the 2026-07
297/// session).
298const COMPANY_NAME_WORDS: &[&str] = &[
299    "company",
300    "legal",
301    "trading",
302    "name",
303    "supplier",
304    "manufacturer",
305    "sponsor",
306    "institution",
307    "journal",
308    "funder",
309    "employer",
310    "law_firm",
311    "affiliation",
312    "reported_by",
313    "buyer",
314    "shipper",
315    "airline",
316];
317
318/// Does `lower` match one of `PERSON_NAME_WORDS`' fragments *other than* the
319/// bare `"name"` catch-all — i.e. is it identifiable as a person-name column
320/// via a specific signal (`given`, `family`, `first_name`, `inspector`...),
321/// not just because its name happens to end in "name"?
322///
323/// Used to keep the `"name"` fragment both lists share (needed to catch
324/// columns like `judge_name`/`director_name` with no more specific word)
325/// from making `given_name`/`company_name` cross-match each other's noise
326/// category — see `is_specifically_company_name` and the `hunt2407.md`
327/// perf-hunt entry this fixes (kyc's `given_name`/`family_name`/
328/// `middle_name` were getting `companies` noise — `drop_legal_form` etc —
329/// applied to them, purely because `COMPANY_NAME_WORDS` also contains
330/// `"name"`).
331fn is_specifically_person_name(lower: &str) -> bool {
332    PERSON_NAME_WORDS
333        .iter()
334        .any(|&w| w != "name" && lower.contains(w))
335}
336
337/// Company-side mirror of `is_specifically_person_name`: matches a
338/// `COMPANY_NAME_WORDS` fragment other than the bare `"name"` catch-all.
339fn is_specifically_company_name(lower: &str) -> bool {
340    COMPANY_NAME_WORDS
341        .iter()
342        .any(|&w| w != "name" && lower.contains(w))
343}
344
345/// Does this noise category ever target a column with this (lowercased) name?
346///
347/// Pure name-pattern predicate shared with `estimate_difficulty` (see
348/// `crate::difficulty`) so the theoretical column-reliability model can never
349/// silently drift from what real generation actually does to each column.
350pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
351    let lower = col_name.to_lowercase();
352    match noise_type {
353        "typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
354        | "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
355        | "blocking_fail" => {
356            // Skip email-like columns — typo/visual noise destroys '@'.
357            // Skip `ip_address` too — it matches the "address" fragment
358            // below by name collision, but it isn't a postal address; this
359            // char-level noise injected letters into octets (BUGS.md C20).
360            if lower.contains("email") || lower.contains("ip_address") {
361                return false;
362            }
363            // Not "phone": that's digit-formatted data, not free text —
364            // char-level typo/homoglyph noise injected letters into phone
365            // numbers (BUGS.md C12). `identifiers`/`corrupt_phone` already
366            // corrupts phone columns digit-aware.
367            contains_any(&lower, &["address", "street", "city"])
368                || contains_any(&lower, PERSON_NAME_WORDS)
369                || contains_any(&lower, COMPANY_NAME_WORDS)
370        }
371        "names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
372            // Exclude columns that are specifically identifiable as a
373            // company/organization name (e.g. `company_name`, `legal_name`)
374            // — the only reason they'd match PERSON_NAME_WORDS at all is the
375            // bare "name" fragment both lists share, not a real person-name
376            // signal. See `is_specifically_company_name`.
377            contains_any(&lower, PERSON_NAME_WORDS) && !is_specifically_company_name(&lower)
378        }
379        "dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
380            contains_any(&lower, &["date", "birth", "incorporation", "founding"])
381        }
382        "missing" | "missing_pattern" => {
383            contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
384        }
385        "identifiers"
386        | "corrupt_email"
387        | "corrupt_phone"
388        | "corrupt_national_id"
389        | "corrupt_siren"
390        | "national_id_corrupt"
391        | "phone_corrupt"
392        | "email_corrupt"
393        | "siren_corrupt" => contains_any(
394            &lower,
395            &[
396                "email",
397                "phone",
398                "ssn",
399                "pan",
400                "passport",
401                "account",
402                "medicare",
403                "license",
404                "siren",
405                "national_id",
406                // Added: `difficulty::base_noise_damage` already models a
407                // 0.6 damage weight for "registration"-matching columns
408                // (e.g. kyc's `registration_number`), but this predicate
409                // never actually targeted them -- a real column with zero
410                // noise coverage regardless of `passes`, floor-ing hell's
411                // f1_max well above its intended ceiling on that schema.
412                "registration",
413            ],
414        ),
415        "extra"
416        | "name_null"
417        | "dob_null"
418        | "blocking_fail_initial"
419        | "blocking_fail_partial"
420        | "fuzzy_match"
421        | "phonetic" => {
422            // Same exclusions as the typo/visual arm above (BUGS.md C12/C20):
423            // skip email (destroys '@') and `ip_address` (not a postal
424            // address despite the name collision), and don't target `phone`
425            // with this char-level noise — `identifiers` already handles it
426            // digit-aware.
427            if lower.contains("email") || lower.contains("ip_address") {
428                return false;
429            }
430            contains_any(&lower, &["address", "street", "city", "note", "comment"])
431                || contains_any(&lower, PERSON_NAME_WORDS)
432                || contains_any(&lower, COMPANY_NAME_WORDS)
433        }
434        "companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
435            // Mirror exclusion: don't apply company-name noise (legal-form
436            // dropping, acronym-ing...) to a column that's specifically
437            // identifiable as a person's name (e.g. `given_name`,
438            // `family_name`) just because it also contains "name". See
439            // `is_specifically_person_name`.
440            contains_any(&lower, COMPANY_NAME_WORDS) && !is_specifically_person_name(&lower)
441        }
442        "addresses" | "address_scramble" | "postal_corrupt" => {
443            // `ip_address` matches "address" by name collision but isn't a
444            // postal address — scrambling/postal-corrupting it doesn't make
445            // sense (see the same exclusion on the typo/visual arm above,
446            // BUGS.md C20). `email_address` is the same collision: it
447            // contains "address" but is an email, not a postal address —
448            // `identifiers`/`corrupt_email` already handles it (perf-hunt
449            // hunt2407.md, confirmed empirically: kyc's `email_address` was
450            // getting postal `address_scramble`/`postal_corrupt` noise).
451            //
452            // "state"/"country" added: `difficulty::base_noise_damage`
453            // already models a 0.5 damage weight for them (same bucket as
454            // address/street/city/postal), but this predicate never
455            // actually targeted them -- e.g. kyc's `residential_state` /
456            // `residential_country` / `registered_country` had zero real
457            // noise coverage regardless of `passes`.
458            !lower.contains("ip_address")
459                && !lower.contains("email")
460                && contains_any(
461                    &lower,
462                    &["address", "street", "postal", "city", "state", "country"],
463                )
464        }
465        "exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
466        | "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
467        | "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
468        | "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
469        | "combo_ultimate" | "french_address" => {
470            // These noise types are handled inline or are no-ops
471            false
472        }
473        // Default: applies to all string columns (except FK columns)
474        _ => true,
475    }
476}
477
478fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
479    let mut matched: Vec<String> = Vec::new();
480    for field in schema.fields() {
481        if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
482            continue;
483        }
484        if exclude_cols.contains(field.name()) {
485            continue;
486        }
487        if noise_type_targets_column(noise_type, field.name()) {
488            matched.push(field.name().clone());
489        }
490    }
491    // For swap, keep exactly 2 columns
492    if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
493        matched.truncate(2);
494    }
495    matched.sort();
496    matched.dedup();
497    matched
498}
499
500fn contains_any(s: &str, patterns: &[&str]) -> bool {
501    patterns.iter().any(|p| s.contains(p))
502}
503
504// ── Noise application ──────────────────────────────────────────────────────
505
506/// Returns the noised batch plus the list of columns it actually targeted
507/// (resolved from the heuristic when `plan_cols` is empty) — callers use
508/// the latter to verify the noise produced a real change, see
509/// `unchanged_row_mask`.
510fn apply_noise_to_batch(
511    rb: &RecordBatch,
512    noise_type: &str,
513    plan_cols: &[String],
514    rng: &mut Rng,
515    exclude_cols: &[String],
516) -> Result<(RecordBatch, Vec<String>), String> {
517    let schema = rb.schema();
518    let n_cols = rb.num_columns();
519
520    // Fallback to heuristic when plan_cols is empty
521    let target_cols: Vec<String> = if plan_cols.is_empty() {
522        match_noise_columns(&schema, noise_type, exclude_cols)
523    } else {
524        plan_cols.to_vec()
525    };
526
527    // Start with None; only clone/modify targeted columns
528    let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();
529
530    let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
531        new_columns[idx]
532            .clone()
533            .unwrap_or_else(|| rb.column(idx).clone())
534    };
535
536    match noise_type {
537        "blocking_fail" => {
538            for col_name in &target_cols {
539                let lower = col_name.to_lowercase();
540                if let Ok(idx) = schema.index_of(col_name) {
541                    let col = get_col(&new_columns, idx, rb);
542                    if lower.contains("first") || lower.contains("given") {
543                        new_columns[idx] =
544                            Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
545                    } else {
546                        new_columns[idx] =
547                            Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
548                    }
549                }
550            }
551        }
552        "swap" | "full_swap" => {
553            if target_cols.len() >= 2 {
554                let name_a = &target_cols[0];
555                let name_b = &target_cols[1];
556                if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
557                    let a = get_col(&new_columns, idx_a, rb);
558                    let b = get_col(&new_columns, idx_b, rb);
559                    new_columns[idx_a] = Some(b);
560                    new_columns[idx_b] = Some(a);
561                }
562            }
563        }
564        "missing_pattern" => {
565            let mut phone_idx: Option<usize> = None;
566            let mut email_idx: Option<usize> = None;
567            for col_name in &target_cols {
568                let lower = col_name.to_lowercase();
569                if let Ok(idx) = schema.index_of(col_name) {
570                    if lower.contains("phone") || lower.contains("mobile") {
571                        phone_idx = Some(idx);
572                    } else if lower.contains("email") {
573                        email_idx = Some(idx);
574                    }
575                }
576            }
577            if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
578                let phone_arr = get_col(&new_columns, pi, rb);
579                let email_arr = get_col(&new_columns, ei, rb);
580                let phone_s = phone_arr.as_string::<i32>();
581                let email_s = email_arr.as_string::<i32>();
582                let n = phone_s.len();
583                let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
584                for i in 0..n {
585                    let pv = if phone_s.is_null(i) {
586                        ""
587                    } else {
588                        phone_s.value(i)
589                    };
590                    let ev = if email_s.is_null(i) {
591                        ""
592                    } else {
593                        email_s.value(i)
594                    };
595                    if !ev.is_empty() && rng.next_usize(2) == 0 {
596                        builder.append_value("");
597                    } else {
598                        builder.append_value(pv);
599                    }
600                }
601                new_columns[pi] = Some(Arc::new(builder.finish()));
602            }
603        }
604        _ => {
605            for col_name in &target_cols {
606                if let Ok(col_idx) = schema.index_of(col_name) {
607                    let col = get_col(&new_columns, col_idx, rb);
608                    let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
609                        .map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
610                    new_columns[col_idx] = Some(result);
611                }
612            }
613        }
614    }
615
616    // Fill remaining from original, unwrap all
617    let final_columns: Vec<ArrayRef> = (0..n_cols)
618        .map(|i| {
619            new_columns[i]
620                .clone()
621                .unwrap_or_else(|| rb.column(i).clone())
622        })
623        .collect();
624
625    let batch =
626        RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))?;
627    Ok((batch, target_cols))
628}
629
630/// Per-row mask (`true` = still unchanged): compares `orig` and `noised` on
631/// every column in `target_cols`, treating null == null as equal. A `true`
632/// row means the noise application was a complete no-op for it — either
633/// because its target columns were already null (a "corrupted" null stays
634/// null) or because the noise function's own internal randomness happened
635/// not to produce a visible variant (several functions in `noise/names.rs`
636/// and `noise/extra.rs` only mutate a row some fraction of the time).
637/// Consumed by the retry loop in the dup-generation closure to guarantee a
638/// duplicate copy that was assigned a noise type actually differs from its
639/// master, instead of silently staying a byte-for-byte duplicate.
640fn unchanged_row_mask(
641    orig: &RecordBatch,
642    noised: &RecordBatch,
643    target_cols: &[String],
644) -> Result<Vec<bool>, String> {
645    let n = orig.num_rows();
646    let mut changed = vec![false; n];
647    let schema = orig.schema();
648    for col_name in target_cols {
649        let Ok(idx) = schema.index_of(col_name) else {
650            continue;
651        };
652        let a = orig.column(idx).as_string::<i32>();
653        let b = noised.column(idx).as_string::<i32>();
654        for (i, is_changed) in changed.iter_mut().enumerate() {
655            if *is_changed {
656                continue;
657            }
658            let av = (!a.is_null(i)).then(|| a.value(i));
659            let bv = (!b.is_null(i)).then(|| b.value(i));
660            if av != bv {
661                *is_changed = true;
662            }
663        }
664    }
665    Ok(changed.into_iter().map(|c| !c).collect())
666}
667
668/// Apply `noise_type` to `orig`, retrying with fresh randomness (merging in
669/// only the still-unchanged rows via `zip`) up to `MAX_NOISE_ATTEMPTS` times,
670/// until every row's target columns differ from `orig` or attempts run out.
671/// Returns the resulting batch, the final per-row "unchanged" mask (`true` =
672/// still byte-identical to `orig` on this noise_type's target columns), and
673/// the `target_cols` this noise_type actually touched.
674///
675/// When `target_cols` comes back empty (no column in this entity's schema
676/// matches `noise_type` at all), every row is returned unchanged (`true`)
677/// with no retries attempted — retrying can't manufacture a matching column
678/// that doesn't exist; this is a schema/config-level gap, not a per-row
679/// randomness one.
680///
681/// The retry merge only `zip()`s columns in `target_cols` (perf-hunt H1,
682/// hunt2407.md): `apply_noise_to_batch` already guarantees every other
683/// column of `retried`/`noisy` is the same `Arc` as `orig`'s (untouched),
684/// so `zip`-ing them is a provably-no-op pass over the full column just to
685/// reconstruct a value that's already sitting there — wasted work that
686/// scales with the entity's *total* column count instead of the handful
687/// this noise_type actually targets. Confirmed empirically: for a 23-column
688/// entity (kyc-sized) vs an 8-column one, this loop alone cost ~1.7x more
689/// wall time despite touching the same single target column either way.
690fn apply_noise_with_retry(
691    orig: &RecordBatch,
692    noise_type: &str,
693    plan_cols: &[String],
694    rng: &mut Rng,
695    exclude_cols: &[String],
696) -> Result<(RecordBatch, Vec<bool>, Vec<String>), String> {
697    const MAX_NOISE_ATTEMPTS: usize = 4;
698    let (mut noisy, target_cols) =
699        apply_noise_to_batch(orig, noise_type, plan_cols, rng, exclude_cols)?;
700    if target_cols.is_empty() {
701        return Ok((noisy, vec![true; orig.num_rows()], target_cols));
702    }
703    let schema = noisy.schema();
704    let target_idxs: Vec<usize> = target_cols
705        .iter()
706        .filter_map(|c| schema.index_of(c).ok())
707        .collect();
708    let mut unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
709    let mut attempt = 1;
710    while unchanged.iter().any(|&u| u) && attempt < MAX_NOISE_ATTEMPTS {
711        let mut retry_rng = Rng::new(rng.next_u64());
712        let (retried, _) =
713            apply_noise_to_batch(orig, noise_type, plan_cols, &mut retry_rng, exclude_cols)?;
714        let mask = arrow::array::BooleanArray::from(unchanged.clone());
715        let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
716        for &i in &target_idxs {
717            merged_cols[i] =
718                arrow::compute::kernels::zip::zip(&mask, retried.column(i), noisy.column(i))
719                    .map_err(|e| format!("zip noise retry col {i}: {e}"))?;
720        }
721        noisy = RecordBatch::try_new(noisy.schema(), merged_cols)
722            .map_err(|e| format!("rebuild retried dup batch: {e}"))?;
723        unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
724        attempt += 1;
725    }
726    Ok((noisy, unchanged, target_cols))
727}
728
729/// Applies one additional noise pass, drawing an independent noise_type per
730/// ROW from `all_types` instead of one type for the entire batch.
731///
732/// The original implementation drew a single `extra_type` for the whole
733/// batch (`all_types[rng.next_usize(all_types.len())]`), applied once. If
734/// that one draw landed on a type with no matching column in this entity's
735/// schema at all (`noise_type_targets_column` false for every column —
736/// common: `dates`/`identifiers`/`addresses`/`companies` are frequently
737/// no-ops for entities that simply have no date/identifier/address/company
738/// column, e.g. an aviation `passenger` row), `apply_noise_with_retry`
739/// returns every row unchanged with *no retry attempted* (retrying can't
740/// manufacture a matching column that doesn't exist) — so the entire batch
741/// silently skipped this pass. With only 2-3 total passes, whether a batch's
742/// "still identical after all passes" fate landed on such a no-op type was
743/// close to a coin flip per pass, and that coin flip's odds shifted with
744/// unrelated RNG-stream position (traced empirically: the same aviation
745/// domain/seed/difficulty produced a real byte-identical-duplicate rate
746/// swinging from ~12% to ~0.003% between two dataset sizes differing by
747/// under 5%, tracking exactly when a per-entity batch boundary was crossed
748/// elsewhere in the run — same root-cause family as `distribute_by_weight`).
749///
750/// Per-row assignment means every batch always has a healthy mix of types,
751/// so a handful of schema-inapplicable categories can no longer determine
752/// the whole batch's fate — same fix class as `noise::apply_random_subtype`,
753/// applied one level up (noise *type* selection, not just sub-type selection
754/// within one already-chosen type).
755fn apply_extra_pass_per_row(
756    noisy: &RecordBatch,
757    unchanged: &mut [bool],
758    all_types: &[&str],
759    rng: &mut Rng,
760    exclude_cols: &[String],
761) -> Result<RecordBatch, String> {
762    let n = noisy.num_rows();
763    if all_types.is_empty() || n == 0 {
764        return Ok(noisy.clone());
765    }
766    let per_row_type: Vec<usize> = (0..n).map(|_| rng.next_usize(all_types.len())).collect();
767
768    // Partition rows by assigned type FIRST (perf-hunt H4, hunt2407.md),
769    // same technique as `noise::apply_random_subtype` — each row needs
770    // noise from exactly one type, but the previous implementation called
771    // `apply_noise_with_retry` on the FULL n-row batch once per active
772    // type and threw away the n - |group| rows that mask didn't select:
773    // O(len(all_types) * n) noise computation instead of O(n). On a schema
774    // where every category has real target columns (e.g. kyc, all 9 hell
775    // categories active — hunt2407.md), that's up to 9x redundant work per
776    // extra pass. Grouping first means each row's noise is computed once,
777    // on a batch sized to just its own group.
778    let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
779    for (i, &t) in per_row_type.iter().enumerate() {
780        groups[t].push(i as u32);
781    }
782
783    let schema = noisy.schema();
784    let mut pass_unchanged = vec![true; n];
785    // column name -> full-length scatter buffer, lazily seeded from
786    // `noisy`'s own values and overwritten in place per group that
787    // touches it. Only columns some active type actually targets get an
788    // entry — untouched columns are left as direct clones of `noisy`'s
789    // arrays at the end, same "don't rebuild what didn't change" principle
790    // as the H1 fix above.
791    let mut touched: HashMap<String, Vec<Option<String>>> = HashMap::new();
792
793    for (type_idx, idxs) in groups.into_iter().enumerate() {
794        if idxs.is_empty() {
795            continue;
796        }
797        let ttype = all_types[type_idx];
798        let idx_arr = UInt64Array::from_iter_values(idxs.iter().map(|&i| i as u64));
799        let subset = pick_rows(noisy, &idx_arr)?;
800        let (noised_subset, subset_unchanged, target_cols) =
801            apply_noise_with_retry(&subset, ttype, &[], rng, exclude_cols)?;
802        for col_name in &target_cols {
803            let Ok(col_idx) = schema.index_of(col_name) else {
804                continue;
805            };
806            let noised_col = noised_subset.column(col_idx).as_string::<i32>();
807            let buf = touched.entry(col_name.clone()).or_insert_with(|| {
808                let orig = noisy.column(col_idx).as_string::<i32>();
809                (0..n)
810                    .map(|i| (!orig.is_null(i)).then(|| orig.value(i).to_string()))
811                    .collect()
812            });
813            for (local_pos, &orig_i) in idxs.iter().enumerate() {
814                buf[orig_i as usize] = (!noised_col.is_null(local_pos))
815                    .then(|| noised_col.value(local_pos).to_string());
816            }
817        }
818        for (local_pos, &orig_i) in idxs.iter().enumerate() {
819            pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
820        }
821    }
822
823    let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
824    for (col_name, buf) in touched {
825        if let Ok(col_idx) = schema.index_of(&col_name) {
826            merged_cols[col_idx] = Arc::new(StringArray::from(buf));
827        }
828    }
829    let current = RecordBatch::try_new(noisy.schema(), merged_cols)
830        .map_err(|e| format!("rebuild extra-pass batch: {e}"))?;
831
832    for (u, pu) in unchanged.iter_mut().zip(pass_unchanged.iter()) {
833        *u = *u && *pu;
834    }
835    Ok(current)
836}
837
838// ── Topological sort for FK dependency ordering ────────────────────────────
839
840/// Order entity plans so that FK target entities are processed before dependents.
841fn topological_sort(
842    plans: &[EntityPlan],
843    name_to_idx: &HashMap<&str, usize>,
844) -> Result<Vec<usize>, String> {
845    let n = plans.len();
846    let mut in_degree = vec![0usize; n];
847    let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
848    for i in 0..n {
849        for remap in &plans[i].fk_remaps {
850            let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
851                continue;
852            };
853            if target == i {
854                continue;
855            }
856            adj[target].push(i); // target → i means target must come before i
857            in_degree[i] += 1;
858        }
859    }
860    let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
861    let mut order = Vec::with_capacity(n);
862    let mut idx = 0;
863    while idx < queue.len() {
864        let v = queue[idx];
865        idx += 1;
866        order.push(v);
867        for &next in &adj[v] {
868            in_degree[next] -= 1;
869            if in_degree[next] == 0 {
870                queue.push(next);
871            }
872        }
873    }
874    if order.len() != n {
875        return Err("FK dependency cycle detected among entities".into());
876    }
877    Ok(order)
878}
879
880/// Dataset sink: writes directly in the requested output format, avoiding
881/// an uncompressed IPC intermediate + full reread when `parquet` is
882/// requested (mirrors the `GtSink` pattern in `gt.rs`).
883pub(crate) enum DatasetWriter {
884    Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
885    Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
886}
887
888impl DatasetWriter {
889    fn new(
890        output_format: &str,
891        path: &str,
892        schema: &Arc<Schema>,
893        metadata: &HashMap<String, String>,
894    ) -> Result<Self, String> {
895        let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
896        if output_format == "parquet" {
897            use parquet::basic::{Compression, ZstdLevel};
898            use parquet::file::properties::WriterProperties;
899            let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
900            let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
901                .iter()
902                .map(|(k, v)| parquet::file::metadata::KeyValue {
903                    key: k.clone(),
904                    value: Some(v.clone()),
905                })
906                .collect();
907            let props = WriterProperties::builder()
908                .set_compression(Compression::ZSTD(zstd))
909                // 1M rows/group (parquet-rs' own default) bounds the
910                // resident, uncompressed column-chunk buffers to one row
911                // group's worth at a time instead of the whole dataset.
912                .set_max_row_group_row_count(Some(1_000_000))
913                .set_key_value_metadata(Some(meta_kv))
914                .build();
915            let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
916                .map_err(|e| format!("ArrowWriter error: {e}"))?;
917            Ok(DatasetWriter::Parquet(Box::new(writer)))
918        } else {
919            let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
920                .map_err(|e| format!("IPC FileWriter error: {e}"))?;
921            Ok(DatasetWriter::Ipc(Box::new(writer)))
922        }
923    }
924
925    pub(crate) fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
926        match self {
927            DatasetWriter::Ipc(w) => w.write(batch).map_err(|e| format!("write ipc: {e}")),
928            DatasetWriter::Parquet(w) => w.write(batch).map_err(|e| format!("write parquet: {e}")),
929        }
930    }
931
932    fn finish(self) -> Result<(), String> {
933        match self {
934            DatasetWriter::Ipc(mut w) => w.finish().map_err(|e| format!("finish ipc: {e}")),
935            DatasetWriter::Parquet(w) => w
936                .close()
937                .map(|_| ())
938                .map_err(|e| format!("close parquet: {e}")),
939        }
940    }
941}
942
943// ── Main pipeline ──────────────────────────────────────────────────────────
944
945pub fn run_pipeline(
946    ctx: &Context,
947    config: &PipelineConfig,
948    output_dir: &str,
949) -> Result<PipelineOutput, String> {
950    run_pipeline_with_progress(ctx, config, output_dir, None)
951}
952
953/// Same as [`run_pipeline`], but invokes `progress(records_written_so_far,
954/// approx_total)` after every base-record batch is written to the dataset
955/// file — `approx_total` is `config.size` (the requested record count; the
956/// actual total including duplicates/HN/canary rows differs slightly, close
957/// enough for a progress display). Callers that don't need progress
958/// reporting should use [`run_pipeline`] instead.
959pub fn run_pipeline_with_progress(
960    ctx: &Context,
961    config: &PipelineConfig,
962    output_dir: &str,
963    mut progress: Option<&mut dyn FnMut(usize, usize)>,
964) -> Result<PipelineOutput, String> {
965    std::fs::create_dir_all(output_dir)
966        .map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
967    let t_start = std::time::Instant::now();
968
969    let t_alloc = t_start.elapsed().as_secs_f64();
970
971    // ── Phase 14: Streaming pipeline ─────────────────────────────────────
972    // Generates, extracts FK/HN, remaps, and sinks each batch per entity
973    // without storing all entity batches in memory.
974    let t1e = std::time::Instant::now();
975
976    // Metadata (including the `dupehell.timestamp` snapshot) is computed
977    // once here and reused for every output file (dataset, graph, GT,
978    // parquet conversions) — computing it per-file let `SystemTime::now()`
979    // drift across a Unix-second boundary on longer runs, making
980    // `dupehell.timestamp` inconsistent between files from the same run.
981    let metadata = build_metadata_map(config);
982    // Build the full schema once (union of all entity columns + metadata)
983    let full_arc = Arc::new(build_full_schema(config, &metadata));
984    let dataset_ext = if config.output_format == "parquet" {
985        "parquet"
986    } else {
987        "ipc"
988    };
989    let dataset_path = format!("{}/{}.{}", output_dir, config.run_id, dataset_ext);
990    let mut writer =
991        DatasetWriter::new(&config.output_format, &dataset_path, &full_arc, &metadata)?;
992
993    // Topological sort for FK dependency processing order
994    let name_to_idx: HashMap<&str, usize> = config
995        .entity_plans
996        .iter()
997        .enumerate()
998        .map(|(i, p)| (p.name.as_str(), i))
999        .collect();
1000    let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;
1001
1002    // Cross-entity storage (lightweight)
1003    let mut fk_pools: FkPoolMap = HashMap::new();
1004    let mut hn_pools: HashMap<String, HnPool> = HashMap::new();
1005
1006    // GT streaming: fed one batch at a time as the dataset is generated
1007    // (see `crate::gt::GtAccumulator`) instead of accumulating full-dataset
1008    // arrays in RAM.
1009    let gt_ext = if config.output_format == "parquet" {
1010        "parquet"
1011    } else {
1012        "ipc"
1013    };
1014    let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
1015    let gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
1016    let mut gt_acc = crate::gt::GtAccumulator::new(&gt_draft_path)?;
1017
1018    let mut global_rid_offset = 0usize;
1019
1020    // ── Graph output (opt-in via --graph) ───────────────────────────────
1021    const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
1022    let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
1023    let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
1024    let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
1025    let mut graph_nodes_path: Option<String> = None;
1026    let mut graph_edges_path: Option<String> = None;
1027    if config.graph_enabled {
1028        let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
1029        let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
1030        node_writer = Some(
1031            crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
1032                .map_err(|e| format!("init node writer: {e}"))?,
1033        );
1034        edge_writer = Some(
1035            crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata)
1036                .map_err(|e| format!("init edge writer: {e}"))?,
1037        );
1038        graph_nodes_path = Some(nodes_ipc);
1039        graph_edges_path = Some(edges_ipc);
1040    }
1041
1042    let mut _t_write = 0.0f64;
1043    let mut _t_meta = 0.0f64;
1044    let mut _t_dup = 0.0f64;
1045    let mut _write_calls = 0u64;
1046    // Compute HN pool size: need at least 2× the max hn_cfg.count for any entity type
1047    let hn_pool_sizing: std::collections::HashMap<&str, usize> =
1048        config
1049            .hard_neg_types
1050            .iter()
1051            .fold(std::collections::HashMap::new(), |mut m, hn| {
1052                let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
1053                *entry = (*entry).max(hn.count);
1054                m
1055            });
1056    let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
1057    // FK remap (fk_remap.rs) samples uniformly at random from this pool via
1058    // `take` — it has no need to see every identifier the entity ever
1059    // produced, only a large-enough sample for the target diversity. Capping
1060    // it (same idea as `hn_max_pool` above) keeps its RAM bounded instead of
1061    // O(n_base) per entity with an identifier column. (`FK_POOL_CAP` is now
1062    // module-level, see its doc comment.)
1063
1064    // Entities whose FK pool is ever read (`fk_pools.get(&remap.target_entity)`
1065    // below, and the identical lookup in `canary::generate_all`, which reuses
1066    // the same `fk_remaps`). Extracting identifiers for an entity outside
1067    // this set is pure waste — the pool built for it is never consulted.
1068    // Entities present only as an FK identifier pool (`--only-entity`),
1069    // never written as graph nodes. A `fk`-typed edge pointing at one of
1070    // these would dangle (target node absent from the nodes output), so
1071    // graph edges are suppressed for remaps into this set below.
1072    let pool_only_names: std::collections::HashSet<&str> = config
1073        .entity_plans
1074        .iter()
1075        .filter(|p| p.pool_only)
1076        .map(|p| p.name.as_str())
1077        .collect();
1078
1079    let fk_targets: std::collections::HashSet<&str> = config
1080        .entity_plans
1081        .iter()
1082        .flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
1083        .collect();
1084
1085    for &plan_idx in &entity_order {
1086        let plan = &config.entity_plans[plan_idx];
1087        if plan.n_base == 0 {
1088            continue;
1089        }
1090
1091        let prefix = entity_prefix(plan_idx);
1092        let col_json_str = &plan.columns_json;
1093
1094        // Per-entity streaming state
1095        let fk_targeted = fk_targets.contains(plan.name.as_str());
1096        let mut fk_builder = arrow::array::StringBuilder::new();
1097        let mut fk_rid_builder = arrow::array::StringBuilder::new();
1098        let mut fk_count: usize = 0;
1099        // Only entities actually referenced by a `hard_neg_types` entry ever
1100        // have their HN pool read (see `hn_pools.get(&hn_cfg.entity_type)`
1101        // below); accumulating slices for the others just pins whole-batch
1102        // Arrow buffers in `hn_pools` for the rest of the run for nothing.
1103        let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
1104        let mut hn_slices: Vec<RecordBatch> = Vec::new();
1105        let mut hn_slice_rids: Vec<Vec<String>> = Vec::new();
1106        let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
1107        let mut fk_rng = Rng::new(config.seed.wrapping_add(42));
1108
1109        // Each noise_type's planned `count` is spread across every batch of
1110        // masters in proportion to batch size (largest-remainder method),
1111        // so duplicates land on masters throughout the whole entity instead
1112        // of being sampled only from whichever batch happens to run last —
1113        // see the "Dups" section below, at the end of each batch iteration.
1114        let batch_bounds: Vec<(usize, usize)> = (0..plan.n_base)
1115            .step_by(crate::entity_gen::BATCH_SIZE)
1116            .map(|offset| {
1117                (
1118                    offset,
1119                    (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE),
1120                )
1121            })
1122            .collect();
1123        let batch_weights: Vec<usize> = batch_bounds.iter().map(|&(_, bn)| bn).collect();
1124        let per_batch_noise_counts: Vec<Vec<usize>> = plan
1125            .noise_types
1126            .iter()
1127            .map(|n| distribute_by_weight(n.count, &batch_weights))
1128            .collect();
1129
1130        // Build col_lookup from first batch (it has the entity's schema)
1131        let mut col_lookup: Option<Vec<Option<usize>>> = None;
1132        // Phase 12.3: null cache per entity plan
1133        let mut null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1134        let mut const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1135
1136        // Stream: generate → FK remap → FK extract → HN sample → IPC write → GT collect
1137        for (batch_idx, offset) in (0..plan.n_base)
1138            .step_by(crate::entity_gen::BATCH_SIZE)
1139            .enumerate()
1140        {
1141            let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
1142            let batch_seed = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);
1143
1144            // Generate batch
1145            let request_json = format!(
1146                r#"{{"entity_name":"{}","n":{},"seed":{},"columns":{},"row_offset":{}}}"#,
1147                plan.name, batch_n, batch_seed, col_json_str, offset,
1148            );
1149            let rb = crate::entity_gen::generate_entity_batch(ctx, &request_json)?;
1150
1151            // FK remap (uses pools from previously processed entities)
1152            let mut fk_edges_by_remap: Vec<(String, ArrayRef, f64)> = Vec::new();
1153            let rb = if !plan.fk_remaps.is_empty() {
1154                let mut r = rb;
1155                for remap in &plan.fk_remaps {
1156                    if let Some(pool) = fk_pools.get(&remap.target_entity) {
1157                        // Suppress edge emission into a pool-only target —
1158                        // its rows were never written as graph nodes, so an
1159                        // edge to them would dangle (see `pool_only_names`).
1160                        let want_edge = config.graph_enabled
1161                            && !pool_only_names.contains(remap.target_entity.as_str());
1162                        let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
1163                            &r,
1164                            pool,
1165                            &remap.source_col,
1166                            &mut fk_rng,
1167                            want_edge,
1168                        )?;
1169                        r = remapped;
1170                        if want_edge && let Some(rids) = target_rids {
1171                            fk_edges_by_remap.push((
1172                                remap.source_col.clone(),
1173                                rids,
1174                                remap.fk_weight,
1175                            ));
1176                        }
1177                    }
1178                }
1179                r
1180            } else {
1181                rb
1182            };
1183
1184            // Extract FK identifiers for this entity's pool, capped at
1185            // FK_POOL_CAP (see comment above).
1186            if fk_targeted
1187                && fk_count < FK_POOL_CAP
1188                && let Some(ref id_col) = plan.identifier_col
1189                && let Some(col) = rb.column_by_name(id_col)
1190            {
1191                let s = col.as_string::<i32>();
1192                for i in 0..s.len() {
1193                    if fk_count >= FK_POOL_CAP {
1194                        break;
1195                    }
1196                    if !s.is_null(i) {
1197                        fk_builder.append_value(s.value(i));
1198                        if config.graph_enabled {
1199                            fk_rid_builder.append_value(record_id_string(global_rid_offset + i));
1200                        }
1201                        fk_count += 1;
1202                    }
1203                }
1204            }
1205
1206            // Pool-only entity (`--only-entity`): its rows exist solely to
1207            // seed `fk_pools` (extracted just above) — no noise, no writing,
1208            // no GT, no graph nodes, and `global_rid_offset` does not
1209            // advance for it (record_ids stay contiguous across the rows
1210            // that ARE written).
1211            if plan.pool_only {
1212                continue;
1213            }
1214
1215            // HN pool: accumulate slices up to max pool size
1216            let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1217            if hn_targeted && hn_accum < hn_max_pool {
1218                let take = (hn_max_pool - hn_accum).min(batch_n);
1219                hn_slices.push(rb.slice(0, take));
1220                // Capture the global record_ids for these sliced rows so the
1221                // HN pool can later resolve `hard_neg` edge targets. Use
1222                // `global_rid_offset`, never the per-entity `offset`.
1223                if config.graph_enabled {
1224                    hn_slice_rids.push(record_id_strs(global_rid_offset..global_rid_offset + take));
1225                }
1226            }
1227
1228            // Build col_lookup from first batch's schema
1229            if col_lookup.is_none() {
1230                let plan_schema = rb.schema();
1231                col_lookup = Some(
1232                    full_arc
1233                        .fields()
1234                        .iter()
1235                        .skip(4)
1236                        .map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1237                        .collect(),
1238                );
1239            }
1240
1241            // Master IDs for this batch
1242            let batch_mids: Vec<String> = (offset..offset + batch_n)
1243                .map(|i| format!("{}-{}", prefix, pad_string(i)))
1244                .collect();
1245            let batch_mids_slice = &batch_mids[..];
1246
1247            // Add metadata + IPC write
1248            let rid_slice = record_id_strs(global_rid_offset..global_rid_offset + batch_n);
1249            let t_m0 = std::time::Instant::now();
1250            let base_rb = add_metadata_and_align(
1251                &rb,
1252                &config.domain,
1253                &plan.name,
1254                &rid_slice,
1255                batch_mids_slice,
1256                &full_arc,
1257                col_lookup.as_ref().unwrap(),
1258                &mut null_cache,
1259                &mut const_arr_cache,
1260            )?;
1261            _t_meta += t_m0.elapsed().as_secs_f64();
1262
1263            // Graph: write base nodes + FK edges (only when --graph)
1264            if config.graph_enabled
1265                && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1266            {
1267                nw.write_batch(&base_rb)
1268                    .map_err(|e| format!("write node: {e}"))?;
1269                let fk_edges_by_remap: Vec<_> = fk_edges_by_remap
1270                    .iter()
1271                    .map(|(subtype, target_rids, weight)| {
1272                        (subtype, target_rids.as_string::<i32>(), *weight)
1273                    })
1274                    .collect();
1275                for (i, rid) in rid_slice.iter().enumerate().take(batch_n) {
1276                    for (subtype, target_rids, weight) in &fk_edges_by_remap {
1277                        ew.push(rid, target_rids.value(i), "fk", subtype, *weight)
1278                            .map_err(|e| format!("push fk edge: {e}"))?;
1279                    }
1280                }
1281            }
1282
1283            let t_w0 = std::time::Instant::now();
1284            writer
1285                .write(&base_rb)
1286                .map_err(|e| format!("write base: {e}"))?;
1287            _t_write += t_w0.elapsed().as_secs_f64();
1288            _write_calls += 1;
1289
1290            // GT collect
1291            gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;
1292
1293            global_rid_offset += batch_n;
1294            if let Some(cb) = &mut progress {
1295                cb(global_rid_offset, config.size);
1296            }
1297
1298            // ── Dups ─────────────────────────────────────────────────────
1299            // Sampled from *this* batch's masters only (`rb`/`batch_n`/
1300            // `offset`), using this batch's share of each noise_type's
1301            // total count (`per_batch_noise_counts`, computed once before
1302            // the loop). Every batch of masters gets duplicates this way,
1303            // instead of concentrating all of them on whichever batch
1304            // happened to run last — see `distribute_by_weight`.
1305            let has_dups: bool = per_batch_noise_counts
1306                .iter()
1307                .any(|counts| counts[batch_idx] > 0);
1308            if has_dups {
1309                let t_d0 = std::time::Instant::now();
1310                let mut dup_batches: Vec<RecordBatch> = Vec::new();
1311                let mut dup_mids_buf: Vec<String> = Vec::new();
1312
1313                // Collect FK columns to exclude from noise
1314                let fk_exclude_cols: Vec<String> = plan
1315                    .fk_remaps
1316                    .iter()
1317                    .map(|r| r.source_col.clone())
1318                    .collect();
1319
1320                // Pre-generate indices + parallel noise (Phase 13c)
1321                let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
1322                    .noise_types
1323                    .iter()
1324                    .enumerate()
1325                    .filter_map(|(ni, n)| {
1326                        let cnt = per_batch_noise_counts[ni][batch_idx];
1327                        if cnt == 0 {
1328                            return None;
1329                        }
1330                        let mut b = UInt64Builder::with_capacity(cnt);
1331                        for _ in 0..cnt {
1332                            b.append_value(batch_rng.next_usize(batch_n) as u64);
1333                        }
1334                        Some((
1335                            b.finish(),
1336                            batch_rng.next_u64(),
1337                            n.noise_type.as_str(),
1338                            n.columns.as_slice(),
1339                            cnt,
1340                        ))
1341                    })
1342                    .collect();
1343
1344                // Distinct active noise_type names for this entity, used by
1345                // additional noise passes (below) to independently draw a type
1346                // to apply on top of the first pass's result.
1347                let all_types: Vec<&str> = plan
1348                    .noise_types
1349                    .iter()
1350                    .map(|n| n.noise_type.as_str())
1351                    .collect();
1352                let all_types_ref = &all_types;
1353                let noise_passes = config.noise_passes.max(1);
1354
1355                type DupNoiseResult = Result<(RecordBatch, Vec<String>, Vec<bool>), String>;
1356                let fk_exclude_ref = &fk_exclude_cols;
1357                let rb_ref = &rb;
1358                let prefix_ref: &str = &prefix;
1359                // `rayon::scope`'s `par_iter().map(...).collect()` runs each
1360                // noise-type bucket on Rayon's already-warm global pool
1361                // instead of spawning a fresh OS thread per bucket per batch
1362                // (the previous `std::thread::scope` pattern) — an isolated
1363                // micro-benchmark (hunt3007/H5, `examples/profile_thread_spawn.rs`)
1364                // measured raw OS thread spawn/join at 12-19x the cost of the
1365                // pooled equivalent at this spawn count/workload shape.
1366                // `par_iter` preserves input order in its output (unlike a
1367                // channel), so `results` ends up in the same `ndata` order as
1368                // before — required downstream since `dup_mids_buf`/
1369                // `global_rid_offset` assignment depends on that order, not on
1370                // which bucket happens to finish first.
1371                let results: Vec<DupNoiseResult> = ndata
1372                    .par_iter()
1373                    .map(|(indices, seed, ntype, cols, cnt)| {
1374                        let dup = pick_rows(rb_ref, indices)?;
1375                        let mut rng = Rng::new(*seed);
1376
1377                        // First pass: the noise_type this bucket was
1378                        // assigned, with the retry-until-changed guarantee
1379                        // (see `apply_noise_with_retry`) so a "noised"
1380                        // duplicate can't silently stay a byte-for-byte
1381                        // copy of its master.
1382                        let (mut noisy, mut unchanged, _) =
1383                            apply_noise_with_retry(&dup, ntype, cols, &mut rng, fk_exclude_ref)?;
1384
1385                        // Additional independent passes (difficulty-
1386                        // controlled, see `schema::DifficultySettings::
1387                        // passes`): each draws its own noise_type from this
1388                        // entity's active list and applies it on top of the
1389                        // previous pass's result, compounding corruption.
1390                        // This — not adding more categories to
1391                        // `noise_types` — is what makes a tier reliably
1392                        // harder without diluting any single category's
1393                        // weight (see the doc comment on `passes`). A row
1394                        // is only "still identical" overall once every
1395                        // pass left it unchanged.
1396                        for _ in 1..noise_passes {
1397                            noisy = apply_extra_pass_per_row(
1398                                &noisy,
1399                                &mut unchanged,
1400                                all_types_ref,
1401                                &mut rng,
1402                                fk_exclude_ref,
1403                            )?;
1404                        }
1405
1406                        let mut mb = Vec::with_capacity(*cnt);
1407                        for j in 0..*cnt {
1408                            // `indices` are local to `rb` (0..batch_n); the master_id
1409                            // is a pure function of (prefix, global index) — same
1410                            // formula as `batch_mids` above — so it's recomputed here
1411                            // instead of being cloned out of a retained master_id_pool.
1412                            let global_idx = offset + indices.value(j) as usize;
1413                            mb.push(format!("{}-{}", prefix_ref, pad_string(global_idx)));
1414                        }
1415                        Ok((noisy, mb, unchanged))
1416                    })
1417                    .collect();
1418                let mut dup_is_identical_buf: Vec<bool> = Vec::new();
1419                for res in results {
1420                    let (rb, mb, unchanged) = res?;
1421                    dup_batches.push(rb);
1422                    dup_mids_buf.extend(mb);
1423                    dup_is_identical_buf.extend(unchanged);
1424                }
1425                _t_dup += t_d0.elapsed().as_secs_f64();
1426
1427                // Write dups
1428                if !dup_batches.is_empty() {
1429                    let dup_total = dup_mids_buf.len();
1430                    let dup_rids = record_id_strs(global_rid_offset..global_rid_offset + dup_total);
1431
1432                    let concated = if dup_batches.len() == 1 {
1433                        dup_batches.into_iter().next().unwrap()
1434                    } else {
1435                        let schema = dup_batches[0].schema();
1436                        let n_fields = dup_batches[0].num_columns();
1437                        let mut concat_arrays = Vec::with_capacity(n_fields);
1438                        for i in 0..n_fields {
1439                            let refs: Vec<&dyn Array> =
1440                                dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
1441                            concat_arrays.push(
1442                                arrow::compute::concat(&refs)
1443                                    .map_err(|e| format!("concat dup col {i}: {e}"))?,
1444                            );
1445                        }
1446                        RecordBatch::try_new(schema, concat_arrays)
1447                            .map_err(|e| format!("concat dups: {e}"))?
1448                    };
1449
1450                    let dup_rb_full = add_metadata_and_align(
1451                        &concated,
1452                        &config.domain,
1453                        &plan.name,
1454                        &dup_rids,
1455                        &dup_mids_buf,
1456                        &full_arc,
1457                        col_lookup.as_ref().unwrap(),
1458                        &mut null_cache,
1459                        &mut const_arr_cache,
1460                    )?;
1461
1462                    let t_wd = std::time::Instant::now();
1463                    writer
1464                        .write(&dup_rb_full)
1465                        .map_err(|e| format!("write dups: {e}"))?;
1466                    _t_write += t_wd.elapsed().as_secs_f64();
1467                    _write_calls += 1;
1468
1469                    // Graph: write duplicate nodes (edges emitted post-GT)
1470                    if config.graph_enabled
1471                        && let Some(nw) = node_writer.as_mut()
1472                    {
1473                        nw.write_batch(&dup_rb_full)
1474                            .map_err(|e| format!("write dup node: {e}"))?;
1475                    }
1476
1477                    let is_identical_arr: ArrayRef =
1478                        Arc::new(arrow::array::BooleanArray::from(dup_is_identical_buf));
1479                    gt_acc.push_dup_batch(
1480                        dup_rb_full.column(0),
1481                        dup_rb_full.column(2),
1482                        dup_rb_full.column(3),
1483                        &is_identical_arr,
1484                    )?;
1485                    global_rid_offset += dup_total;
1486                }
1487            }
1488        }
1489
1490        // Save FK pool for cross-entity remapping
1491        if let Some(id_col) = &plan.identifier_col
1492            && fk_builder.len() > 0
1493        {
1494            let arr = Arc::new(fk_builder.finish()) as ArrayRef;
1495            if config.graph_enabled {
1496                // Graph mode: also persist the target record_ids (column 1)
1497                // so FK edges can be emitted by later-source entities. Same
1498                // bounds as the identifier column above (FK_POOL_CAP gated).
1499                let rid_arr = Arc::new(fk_rid_builder.finish()) as ArrayRef;
1500                let schema = Arc::new(Schema::new(vec![
1501                    Field::new(id_col, DataType::Utf8, true),
1502                    Field::new("record_id", DataType::Utf8, true),
1503                ]));
1504                if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
1505                    fk_pools.insert(plan.name.clone(), pool_rb);
1506                }
1507            } else {
1508                let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
1509                if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
1510                    fk_pools.insert(plan.name.clone(), pool_rb);
1511                }
1512            }
1513        }
1514
1515        // Build HN pool from collected slices
1516        if !hn_slices.is_empty() {
1517            let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1518            let batch = if hn_slices.len() == 1 {
1519                // `rb.slice(0, take)` (pushed above) is zero-copy: it keeps
1520                // the *entire* source batch's buffers alive (up to
1521                // BATCH_SIZE=500_000 rows × all columns) for the rest of the
1522                // run just to serve `take` rows. Re-materializing through the
1523                // `take` kernel (same helper as `pick_rows` below) builds
1524                // compact, right-sized buffers instead.
1525                let only = hn_slices.into_iter().next().unwrap();
1526                let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
1527                pick_rows(&only, &idx)?
1528            } else {
1529                let schema = hn_slices[0].schema();
1530                let n_fields = schema.fields().len();
1531                let mut concat_arrays = Vec::with_capacity(n_fields);
1532                for i in 0..n_fields {
1533                    let refs: Vec<&dyn Array> =
1534                        hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
1535                    concat_arrays.push(
1536                        arrow::compute::concat(&refs)
1537                            .map_err(|e| format!("hn pool concat col {i}: {e}"))?,
1538                    );
1539                }
1540                RecordBatch::try_new(schema, concat_arrays)
1541                    .map_err(|e| format!("hn pool concat batch: {e}"))?
1542            };
1543            // `record_ids` aligns positionally with the concatenated pool rows.
1544            let record_ids = if config.graph_enabled {
1545                hn_slice_rids.into_iter().flatten().collect::<Vec<String>>()
1546            } else {
1547                Vec::new()
1548            };
1549            hn_pools.insert(
1550                plan.name.clone(),
1551                HnPool {
1552                    batch,
1553                    total_count,
1554                    record_ids,
1555                },
1556            );
1557        }
1558    }
1559    let t1e_elapsed = t1e.elapsed().as_secs_f64();
1560    log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
1561
1562    // ── Phase 2: Hard negatives ────────────────────────────────────────────
1563    let t2 = std::time::Instant::now();
1564    // Phase 12.3: separate null cache for HN section (different entity types)
1565    let mut hn_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1566    let mut hn_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1567    // Global counter for HN master_ids (shared across every hard_neg_types
1568    // entry in this run) instead of a random draw from a fixed 10M space:
1569    // at tens/hundreds of thousands of HN rows, `next_usize(10_000_000)`
1570    // has a real birthday-paradox chance of two unrelated HN rows landing
1571    // on the same master_id, which `gt.rs` would then count as sharing a
1572    // master — silently mislabeling a hard negative as a match. A counter
1573    // is unique by construction, eliminating the collision entirely.
1574    let mut hn_master_id_counter: u64 = 0;
1575    for hn_cfg in config.hard_neg_types.iter() {
1576        if hn_cfg.count == 0 {
1577            continue;
1578        }
1579        let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
1580            continue;
1581        };
1582        let hn_total = pool_data.total_count;
1583
1584        let pool_rb = &pool_data.batch;
1585        let n_hn_max = hn_cfg.count.min(hn_total / 2);
1586        if n_hn_max == 0 {
1587            continue;
1588        }
1589
1590        let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
1591            pool_rb,
1592            &hn_cfg.config_json,
1593            n_hn_max,
1594            config.seed.wrapping_add(200),
1595            config.graph_enabled,
1596        )?;
1597
1598        let n_hn = hn_rb.num_rows();
1599        if n_hn == 0 {
1600            continue;
1601        }
1602
1603        let hn_rids = record_id_strs(global_rid_offset..global_rid_offset + n_hn);
1604
1605        let hn_mids: Vec<String> = (0..n_hn)
1606            .map(|_i| {
1607                let id = hn_master_id_counter;
1608                hn_master_id_counter += 1;
1609                format!("HN-{:09}", id)
1610            })
1611            .collect();
1612
1613        // Build col lookup for the HN entity type schema
1614        let hn_schema = hn_rb.schema();
1615        let hn_col_lookup: Vec<Option<usize>> = full_arc
1616            .fields()
1617            .iter()
1618            .skip(4)
1619            .map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1620            .collect();
1621
1622        let hn_rb_full = add_metadata_and_align(
1623            &hn_rb,
1624            &config.domain,
1625            &hn_cfg.entity_type,
1626            &hn_rids,
1627            &hn_mids,
1628            &full_arc,
1629            &hn_col_lookup,
1630            &mut hn_null_cache,
1631            &mut hn_const_arr_cache,
1632        )?;
1633
1634        let t_wh = std::time::Instant::now();
1635        writer
1636            .write(&hn_rb_full)
1637            .map_err(|e| format!("write hn: {e}"))?;
1638        _t_write += t_wh.elapsed().as_secs_f64();
1639        _write_calls += 1;
1640
1641        // Graph: write HN nodes + hard_neg edges (source = HN rid,
1642        // target = the pool row referenced by idx_a)
1643        if config.graph_enabled
1644            && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1645        {
1646            nw.write_batch(&hn_rb_full)
1647                .map_err(|e| format!("write hn node: {e}"))?;
1648            if let Some((idx_a, pattern)) = hn_src {
1649                for i in 0..n_hn {
1650                    let tgt = &pool_data.record_ids[idx_a[i]];
1651                    ew.push(&hn_rids[i], tgt, "hard_neg", &pattern, 1.0)
1652                        .map_err(|e| format!("push hn edge: {e}"))?;
1653                }
1654            }
1655        }
1656
1657        // Phase 13: collect ArrayRefs instead of per-row StringBuilder
1658        gt_acc.push_other_batch(
1659            hn_rb_full.column(0),
1660            hn_rb_full.column(2),
1661            hn_rb_full.column(3),
1662        )?;
1663
1664        global_rid_offset += n_hn;
1665    }
1666    let t2_elapsed = t2.elapsed().as_secs_f64();
1667
1668    // ── Canary records ────────────────────────────────────────────────────
1669    {
1670        let mut canary_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1671        let mut canary_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1672        crate::canary::generate_all(
1673            ctx,
1674            config,
1675            &full_arc,
1676            &mut canary_null_cache,
1677            &mut canary_const_arr_cache,
1678            &mut global_rid_offset,
1679            &fk_pools,
1680            &mut writer,
1681            &mut node_writer,
1682            &mut gt_acc,
1683        )?;
1684    }
1685
1686    // ── Phase 3: Finalize + GT (streamed, see crate::gt::GtAccumulator) ────
1687    let t3 = std::time::Instant::now();
1688    writer
1689        .finish()
1690        .map_err(|e| format!("finish dataset writer: {e}"))?;
1691    let t3a_elapsed = t3.elapsed().as_secs_f64();
1692
1693    let t3b = std::time::Instant::now();
1694    let t_gt0 = std::time::Instant::now();
1695    let crate::gt::GtResult {
1696        n_exact_dup,
1697        n_fuzzy_dup,
1698        n_hard_neg,
1699        n_unique,
1700        n_masters,
1701        cluster_map,
1702    } = gt_acc.finish(
1703        config.difficulty.as_str(),
1704        &config.output_format,
1705        &gt_path,
1706        &metadata,
1707        config.graph_enabled,
1708    )?;
1709    let _t_gt_compute = 0.0f64;
1710    let _t_gt_write = t_gt0.elapsed().as_secs_f64();
1711    let t3b_elapsed = t3b.elapsed().as_secs_f64();
1712    let _t3_elapsed = t3.elapsed().as_secs_f64();
1713
1714    // ── Graph: emit duplicate-cluster edges from the post-GT cluster_map ──
1715    if config.graph_enabled
1716        && let Some(ew) = edge_writer.as_mut()
1717    {
1718        crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
1719            .map_err(|e| format!("push dup clusters: {e}"))?;
1720    }
1721
1722    let duration = t_start.elapsed().as_secs_f64();
1723
1724    log::debug!(
1725        "[pipeline_timing] alloc={:.3}s  sink={:.3}s  hn={:.3}s  merge={:.3}s  meta={:.3}s/dup={:.3}s/write={:.3}s  gt(comp={:.3}s+write={:.3}s)={:.3}s  total={:.3}s",
1726        t_alloc,
1727        t1e_elapsed,
1728        t2_elapsed,
1729        t3a_elapsed,
1730        _t_meta,
1731        _t_dup,
1732        _t_write,
1733        _t_gt_compute,
1734        _t_gt_write,
1735        t3b_elapsed,
1736        duration,
1737    );
1738
1739    let stats = PipelineStats {
1740        total_records: global_rid_offset,
1741        exact_dups: n_exact_dup,
1742        fuzzy_dups: n_fuzzy_dup,
1743        hard_negs: n_hard_neg,
1744        uniques: n_unique,
1745        masters: n_masters,
1746    };
1747
1748    // ── Finalize graph writers (base nodes + FK edges emitted so far) ──
1749    let mut graph_nodes_final: Option<String> = None;
1750    let mut graph_edges_final: Option<String> = None;
1751    if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
1752        nw.finish()
1753            .map_err(|e| format!("finish node writer: {e}"))?;
1754        ew.finish()
1755            .map_err(|e| format!("finish edge writer: {e}"))?;
1756        let mut np = graph_nodes_path.clone().unwrap();
1757        let mut ep = graph_edges_path.clone().unwrap();
1758        if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
1759            let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
1760            let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
1761            convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
1762            convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
1763            std::fs::remove_file(&np).ok();
1764            std::fs::remove_file(&ep).ok();
1765            np = np_pq;
1766            ep = ep_pq;
1767        }
1768        graph_nodes_final = Some(np);
1769        graph_edges_final = Some(ep);
1770    }
1771
1772    Ok(PipelineOutput {
1773        output_files: vec![dataset_path],
1774        gt_file: gt_path,
1775        stats,
1776        nodes: graph_nodes_final,
1777        edges: graph_edges_final,
1778    })
1779}
1780
1781/// IPC → Parquet (ZSTD) conversion for graph files, mirroring the dataset
1782/// conversion above. The node IPC already carries the `dupehell.*` metadata,
1783/// which is forwarded to the Parquet key/value metadata.
1784fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
1785    use arrow::ipc::reader::FileReader;
1786
1787    let ipc_file =
1788        std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
1789    let ipc_reader =
1790        FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
1791    let schema = ipc_reader.schema();
1792
1793    let parquet_file =
1794        std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
1795    let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
1796    let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
1797        .metadata()
1798        .iter()
1799        .map(|(k, v)| parquet::file::metadata::KeyValue {
1800            key: k.clone(),
1801            value: Some(v.clone()),
1802        })
1803        .collect();
1804    let props = parquet::file::properties::WriterProperties::builder()
1805        .set_compression(parquet::basic::Compression::ZSTD(zstd_level))
1806        .set_max_row_group_row_count(Some(1_000_000))
1807        .set_key_value_metadata(Some(meta_kv))
1808        .build();
1809    let mut parquet_writer =
1810        parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
1811            .map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;
1812
1813    for batch_result in ipc_reader {
1814        let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
1815        parquet_writer
1816            .write(&batch)
1817            .map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
1818    }
1819    parquet_writer
1820        .close()
1821        .map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
1822    Ok(())
1823}
1824
1825// ── Metadata injection ──────────────────────────────────────────────────────
1826
1827fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
1828    let ts = std::time::SystemTime::now()
1829        .duration_since(std::time::UNIX_EPOCH)
1830        .map(|d| d.as_secs().to_string())
1831        .unwrap_or_else(|_| "unknown".into());
1832    HashMap::from([
1833        (
1834            "dupehell.generator".into(),
1835            format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
1836        ),
1837        (
1838            "dupehell.provenance".into(),
1839            "dupehell-synthetic-data".into(),
1840        ),
1841        ("dupehell.license".into(), "MIT".into()),
1842        (
1843            "dupehell.purpose".into(),
1844            "Educational Use Only — Record Linkage Benchmarking".into(),
1845        ),
1846        (
1847            "dupehell.url".into(),
1848            "https://github.com/vntoinekaio/DupeHell".into(),
1849        ),
1850        ("dupehell.domain".into(), config.domain.clone()),
1851        ("dupehell.size".into(), config.size.to_string()),
1852        ("dupehell.seed".into(), config.seed.to_string()),
1853        ("dupehell.difficulty".into(), config.difficulty.clone()),
1854        ("dupehell.run_id".into(), config.run_id.clone()),
1855        ("dupehell.timestamp".into(), ts),
1856        (
1857            "dupehell.structural_key_columns".into(),
1858            structural_key_columns_note(config),
1859        ),
1860    ])
1861}
1862
1863/// Lists the `{entity}_id` columns that are structural join keys (linking an
1864/// entity's rows to its flattened child tables), not attributes to feed an
1865/// ER model — they never receive noise and stay identical across all
1866/// duplicates of the same master_id, unlike `record_id`.
1867fn structural_key_columns_note(config: &PipelineConfig) -> String {
1868    let cols: Vec<&str> = config
1869        .entity_plans
1870        .iter()
1871        .filter_map(|p| p.identifier_col.as_deref())
1872        .collect();
1873    format!(
1874        "{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
1875        cols.join(", ")
1876    )
1877}
1878
1879// ── Schema alignment ───────────────────────────────────────────────────────
1880
1881/// Build the union schema from all entity plans (all columns + metadata).
1882fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
1883    let mut field_map: Vec<(String, DataType, bool)> = Vec::new();
1884    let metadata_fields = ["record_id", "domain", "entity_type", "master_id"];
1885    for mf in &metadata_fields {
1886        field_map.push((mf.to_string(), DataType::Utf8, false));
1887    }
1888    for plan in &config.entity_plans {
1889        // Pool-only entities (`--only-entity`) are never written to output —
1890        // excluding their columns here is what gives `--only-entity` its
1891        // narrow, single-entity schema for free, with no separate filtering
1892        // step downstream.
1893        if plan.pool_only {
1894            continue;
1895        }
1896        let cols: Vec<serde_json::Value> =
1897            serde_json::from_str(&plan.columns_json).unwrap_or_default();
1898        for col in &cols {
1899            let name = col["name"].as_str().unwrap_or("").to_string();
1900            if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
1901                continue;
1902            }
1903            let col_type = col_type_from_request(col);
1904            field_map.push((name, col_type, true));
1905        }
1906    }
1907    let fields: Vec<Field> = field_map
1908        .into_iter()
1909        .map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
1910        .collect();
1911    Schema::new(fields).with_metadata(metadata.clone())
1912}
1913
1914/// Map a column's JSON type string to Arrow DataType.
1915fn col_type_from_request(v: &serde_json::Value) -> DataType {
1916    match v["type"].as_str() {
1917        Some("int") => DataType::Int64,
1918        Some("float") => DataType::Float64,
1919        Some("boolean") => DataType::Boolean,
1920        Some("date") | Some("datetime") => DataType::Utf8,
1921        _ => DataType::Utf8,
1922    }
1923}
1924
1925// ── Helpers ─────────────────────────────────────────────────────────────────
1926
1927#[allow(clippy::too_many_arguments)]
1928pub(crate) fn add_metadata_and_align(
1929    rb: &RecordBatch,
1930    domain: &str,
1931    entity_type: &str,
1932    record_ids: &[String],
1933    master_ids: &[String],
1934    full_arc: &Arc<Schema>,
1935    col_lookup: &[Option<usize>],
1936    null_cache: &mut FxHashMap<(DataType, usize), ArrayRef>,
1937    const_arr_cache: &mut FxHashMap<(String, usize), ArrayRef>,
1938) -> Result<RecordBatch, String> {
1939    let n = rb.num_rows();
1940    // `domain` is constant for the whole run and `entity_type` repeats
1941    // across every batch of a given entity (and often across HN/canary
1942    // sections too) — cache the built repeat-array per (value, n) instead
1943    // of rebuilding an n-length StringArray on every call, same idea as
1944    // `null_cache` above for the always-null columns.
1945    let domain_arr = const_arr_cache
1946        .entry((domain.to_string(), n))
1947        .or_insert_with(|| {
1948            Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
1949                domain, n,
1950            ))) as ArrayRef
1951        })
1952        .clone();
1953    let et_arr = const_arr_cache
1954        .entry((entity_type.to_string(), n))
1955        .or_insert_with(|| {
1956            Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
1957                entity_type,
1958                n,
1959            ))) as ArrayRef
1960        })
1961        .clone();
1962    let rid_arr = Arc::new(StringArray::from_iter_values(
1963        record_ids.iter().map(|s| s.as_str()),
1964    )) as ArrayRef;
1965    let mid_arr = Arc::new(StringArray::from_iter_values(
1966        master_ids.iter().map(|s| s.as_str()),
1967    )) as ArrayRef;
1968
1969    let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
1970    all_arrays.push(rid_arr);
1971    all_arrays.push(domain_arr);
1972    all_arrays.push(et_arr);
1973    all_arrays.push(mid_arr);
1974
1975    // P4: Vec<Option<usize>> eliminates HashMap lookup per field (just index)
1976    // Phase 12.3: cache null arrays per (DataType, n) — avoids new_null_array overhead per batch
1977    for (i, maybe_idx) in col_lookup.iter().enumerate() {
1978        match maybe_idx {
1979            Some(idx) => all_arrays.push(rb.column(*idx).clone()),
1980            None => {
1981                let dt = full_arc.field(i + 4).data_type();
1982                let arr = null_cache
1983                    .entry((dt.clone(), n))
1984                    .or_insert_with(|| arrow::array::new_null_array(dt, n));
1985                all_arrays.push(Arc::clone(arr));
1986            }
1987        }
1988    }
1989
1990    // P3: Arc::clone(full_arc) avoids full schema clone (just atomic increment)
1991    RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
1992        .map_err(|e| format!("add_metadata_and_align RecordBatch: {e}"))
1993}
1994
1995fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
1996    use arrow::compute::take;
1997    let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
1998        .map(|i| take(rb.column(i), indices, None))
1999        .collect::<Result<Vec<_>, _>>()
2000        .map_err(|e| format!("take error: {e}"))?;
2001
2002    RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
2003}
2004
2005#[cfg(test)]
2006mod tests {
2007    use super::*;
2008    use crate::schema::{build_pipeline_config, load_schema};
2009
2010    /// Regression for the perf-hunt cross-contamination bug (hunt2407.md):
2011    /// `COMPANY_NAME_WORDS` and `PERSON_NAME_WORDS` both contain the bare
2012    /// `"name"` fragment, which used to let `given_name`/`family_name`/
2013    /// `middle_name` match the `companies` noise category (and would
2014    /// symmetrically let `company_name`/`legal_name` match `names`), and let
2015    /// `email_address` match the `addresses` category via `"address"`.
2016    #[test]
2017    fn test_no_person_company_name_cross_contamination() {
2018        for col in [
2019            "given_name",
2020            "family_name",
2021            "middle_name",
2022            "first_name",
2023            "last_name",
2024        ] {
2025            assert!(
2026                noise_type_targets_column("names", col),
2027                "{col} should still match 'names'"
2028            );
2029            assert!(
2030                !noise_type_targets_column("companies", col),
2031                "{col} should NOT match 'companies' (person name, not a company name)"
2032            );
2033        }
2034        for col in ["company_name", "legal_name", "trading_name"] {
2035            assert!(
2036                noise_type_targets_column("companies", col),
2037                "{col} should still match 'companies'"
2038            );
2039            assert!(
2040                !noise_type_targets_column("names", col),
2041                "{col} should NOT match 'names' (company name, not a person name)"
2042            );
2043        }
2044        assert!(
2045            !noise_type_targets_column("addresses", "email_address"),
2046            "email_address should NOT match 'addresses' (it's an email, not a postal address)"
2047        );
2048        assert!(
2049            noise_type_targets_column("identifiers", "email_address"),
2050            "email_address should still match 'identifiers' (corrupt_email)"
2051        );
2052        assert!(
2053            noise_type_targets_column("addresses", "residential_address_line1"),
2054            "a real postal address column should still match 'addresses'"
2055        );
2056    }
2057
2058    fn manifest_dir() -> std::path::PathBuf {
2059        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).to_path_buf()
2060    }
2061
2062    /// Unique per-test scratch dir under `target/`, so parallel test runs
2063    /// never collide on output filenames and cleanup is trivial.
2064    fn scratch_dir(name: &str) -> std::path::PathBuf {
2065        let dir = manifest_dir()
2066            .join("target")
2067            .join("pipeline_test_tmp")
2068            .join(format!("{name}_{}", std::process::id()));
2069        let _ = std::fs::remove_dir_all(&dir);
2070        std::fs::create_dir_all(&dir).expect("create scratch dir");
2071        dir
2072    }
2073
2074    fn run_kyc(
2075        output_dir: &std::path::Path,
2076        output_format: &str,
2077        graph_enabled: bool,
2078    ) -> PipelineOutput {
2079        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2080        let mut ctx = Context::new(
2081            "kyc",
2082            "en",
2083            &manifest_dir().join("assets/pools").to_string_lossy(),
2084        )
2085        .expect("load context");
2086        let config = build_pipeline_config(
2087            "kyc",
2088            200,
2089            42,
2090            "medium",
2091            0.1,
2092            0.3,
2093            &schema,
2094            "pipeline_test_run",
2095            output_format,
2096            graph_enabled,
2097            "parquet",
2098            None,
2099        )
2100        .expect("build config");
2101        ctx.enable_watermark(&config.domain, config.size, config.seed);
2102        run_pipeline(&ctx, &config, &output_dir.to_string_lossy()).expect("run_pipeline")
2103    }
2104
2105    #[test]
2106    fn test_run_pipeline_produces_dataset_and_gt() {
2107        let dir = scratch_dir("basic");
2108        let output = run_kyc(&dir, "parquet", false);
2109
2110        assert_eq!(output.output_files.len(), 1);
2111        for f in &output.output_files {
2112            assert!(std::path::Path::new(f).exists(), "missing dataset file {f}");
2113        }
2114        assert!(std::path::Path::new(&output.gt_file).exists());
2115        assert!(output.nodes.is_none());
2116        assert!(output.edges.is_none());
2117
2118        // Stats must be internally consistent: exact_dup/fuzzy_dup/hard_neg/
2119        // unique are a partition of the classified records, which is <=
2120        // total_records (total_records also includes unclassified canary
2121        // watermark rows).
2122        let s = &output.stats;
2123        assert!(s.total_records >= s.exact_dups + s.fuzzy_dups + s.hard_negs + s.uniques);
2124        assert!(s.total_records > 0);
2125        assert!(s.masters > 0);
2126
2127        let _ = std::fs::remove_dir_all(&dir);
2128    }
2129
2130    #[test]
2131    fn test_run_pipeline_ipc_format() {
2132        let dir = scratch_dir("ipc");
2133        let output = run_kyc(&dir, "ipc", false);
2134        assert!(output.output_files[0].ends_with(".ipc"));
2135        assert!(std::path::Path::new(&output.output_files[0]).exists());
2136        let _ = std::fs::remove_dir_all(&dir);
2137    }
2138
2139    #[test]
2140    fn test_run_pipeline_graph_output() {
2141        let dir = scratch_dir("graph");
2142        let output = run_kyc(&dir, "parquet", true);
2143        let nodes = output.nodes.expect("nodes path");
2144        let edges = output.edges.expect("edges path");
2145        assert!(std::path::Path::new(&nodes).exists());
2146        assert!(std::path::Path::new(&edges).exists());
2147        let _ = std::fs::remove_dir_all(&dir);
2148    }
2149
2150    #[test]
2151    fn test_fk_edge_weight_from_schema_and_default() {
2152        // aviation's flight->aircraft fk_remap declares fk_weight=0.5
2153        // (schemas/aviation.json); flight->airline leaves the field at
2154        // its schema default of 1.0. Confirms `FkRemapEntry::fk_weight`
2155        // is threaded through to the `weight` column of `fk`-typed graph
2156        // edges instead of the old hardcoded 1.0 constant, and that
2157        // domains/relations without the field still default correctly.
2158        let dir = scratch_dir("fk_weight");
2159        let schema =
2160            load_schema("aviation", &manifest_dir().join("schemas")).expect("load aviation.json");
2161        let mut ctx = Context::new(
2162            "aviation",
2163            "en",
2164            &manifest_dir().join("assets/pools").to_string_lossy(),
2165        )
2166        .expect("load context");
2167        let config = build_pipeline_config(
2168            "aviation",
2169            300,
2170            42,
2171            "medium",
2172            0.1,
2173            0.3,
2174            &schema,
2175            "pipeline_test_fk_weight",
2176            "ipc",
2177            true,
2178            "ipc",
2179            None,
2180        )
2181        .expect("build config");
2182        ctx.enable_watermark(&config.domain, config.size, config.seed);
2183        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2184
2185        let edges_path = output.edges.expect("edges path");
2186        let file = std::fs::File::open(&edges_path).expect("open edges file");
2187        let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2188
2189        let mut subtypes_seen: std::collections::HashMap<String, f64> =
2190            std::collections::HashMap::new();
2191        for batch in reader {
2192            let batch = batch.expect("read edges batch");
2193            let edge_type = batch
2194                .column_by_name("edge_type")
2195                .expect("edge_type col")
2196                .as_any()
2197                .downcast_ref::<arrow::array::StringArray>()
2198                .expect("edge_type as string");
2199            let subtype = batch
2200                .column_by_name("subtype")
2201                .expect("subtype col")
2202                .as_any()
2203                .downcast_ref::<arrow::array::StringArray>()
2204                .expect("subtype as string");
2205            let weight = batch
2206                .column_by_name("weight")
2207                .expect("weight col")
2208                .as_any()
2209                .downcast_ref::<arrow::array::Float64Array>()
2210                .expect("weight as f64");
2211            for i in 0..batch.num_rows() {
2212                if edge_type.value(i) == "fk" {
2213                    subtypes_seen.insert(subtype.value(i).to_string(), weight.value(i));
2214                }
2215            }
2216        }
2217
2218        assert_eq!(
2219            subtypes_seen.get("aircraft_id"),
2220            Some(&0.5),
2221            "flight->aircraft fk_weight from schema not propagated to edge weight"
2222        );
2223        assert_eq!(
2224            subtypes_seen.get("airline_id"),
2225            Some(&1.0),
2226            "flight->airline fk_weight (schema default) not propagated to edge weight"
2227        );
2228
2229        let _ = std::fs::remove_dir_all(&dir);
2230    }
2231
2232    /// End-to-end regression for the `ClusterCsr` rewrite of `cluster_map`
2233    /// (`project_csr_cluster_map_backlog`): drives the real pipeline
2234    /// (`gt::GtAccumulator::finish` -> `ClusterCsr::build` ->
2235    /// `graph_gen::push_dup_clusters`) instead of testing either module in
2236    /// isolation, and checks the emitted `exact_dup`/`fuzzy_dup` edges are
2237    /// well-formed and consistent with the run's own GT stats.
2238    #[test]
2239    fn test_run_pipeline_dup_cluster_edges_well_formed() {
2240        let dir = scratch_dir("dup_cluster_edges");
2241        // `run_kyc` always requests parquet-format graph output regardless
2242        // of its `output_format` arg, so build the config directly here to
2243        // get IPC edges (matches `test_fk_edge_weight_from_schema_and_default`'s
2244        // approach for the same reason).
2245        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2246        let mut ctx = Context::new(
2247            "kyc",
2248            "en",
2249            &manifest_dir().join("assets/pools").to_string_lossy(),
2250        )
2251        .expect("load context");
2252        let config = build_pipeline_config(
2253            "kyc",
2254            200,
2255            42,
2256            "medium",
2257            0.1,
2258            0.3,
2259            &schema,
2260            "pipeline_test_dup_cluster_edges",
2261            "ipc",
2262            true,
2263            "ipc",
2264            None,
2265        )
2266        .expect("build config");
2267        ctx.enable_watermark(&config.domain, config.size, config.seed);
2268        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2269
2270        let edges_path = output.edges.expect("edges path");
2271        let file = std::fs::File::open(&edges_path).expect("open edges file");
2272        let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2273
2274        let mut n_exact = 0usize;
2275        let mut n_fuzzy = 0usize;
2276        for batch in reader {
2277            let batch = batch.expect("read edges batch");
2278            let src = batch
2279                .column_by_name("source_node_id")
2280                .expect("source_node_id col")
2281                .as_any()
2282                .downcast_ref::<arrow::array::StringArray>()
2283                .expect("source_node_id as string");
2284            let tgt = batch
2285                .column_by_name("target_node_id")
2286                .expect("target_node_id col")
2287                .as_any()
2288                .downcast_ref::<arrow::array::StringArray>()
2289                .expect("target_node_id as string");
2290            let edge_type = batch
2291                .column_by_name("edge_type")
2292                .expect("edge_type col")
2293                .as_any()
2294                .downcast_ref::<arrow::array::StringArray>()
2295                .expect("edge_type as string");
2296            let weight = batch
2297                .column_by_name("weight")
2298                .expect("weight col")
2299                .as_any()
2300                .downcast_ref::<arrow::array::Float64Array>()
2301                .expect("weight as f64");
2302
2303            for i in 0..batch.num_rows() {
2304                match edge_type.value(i) {
2305                    "exact_dup" => n_exact += 1,
2306                    "fuzzy_dup" => n_fuzzy += 1,
2307                    _ => continue,
2308                }
2309                // Every endpoint must round-trip through `parse_record_idx`
2310                // -- i.e. `ClusterCsr`'s packed indices were correctly
2311                // unpacked back into real record_ids by `push_dup_clusters`,
2312                // not garbage or truncated.
2313                assert!(
2314                    parse_record_idx(src.value(i)).is_some(),
2315                    "malformed source_node_id {:?} on a dup-cluster edge",
2316                    src.value(i)
2317                );
2318                assert!(
2319                    parse_record_idx(tgt.value(i)).is_some(),
2320                    "malformed target_node_id {:?} on a dup-cluster edge",
2321                    tgt.value(i)
2322                );
2323                assert!((weight.value(i) - 1.0).abs() < 1e-9);
2324            }
2325        }
2326
2327        // kyc/medium/seed 42/size 200 always produces some duplicates (see
2328        // `output.stats.exact_dups`/`fuzzy_dups`), so the cluster-edge pass
2329        // must actually have emitted something -- an empty result here would
2330        // mean `ClusterCsr` silently dropped every cluster.
2331        assert!(n_exact + n_fuzzy > 0, "no dup-cluster edges emitted");
2332        assert!(output.stats.exact_dups + output.stats.fuzzy_dups > 0);
2333
2334        let _ = std::fs::remove_dir_all(&dir);
2335    }
2336
2337    #[test]
2338    fn test_run_pipeline_deterministic_stats() {
2339        let dir_a = scratch_dir("det_a");
2340        let dir_b = scratch_dir("det_b");
2341        let a = run_kyc(&dir_a, "parquet", false);
2342        let b = run_kyc(&dir_b, "parquet", false);
2343        assert_eq!(a.stats.total_records, b.stats.total_records);
2344        assert_eq!(a.stats.exact_dups, b.stats.exact_dups);
2345        assert_eq!(a.stats.fuzzy_dups, b.stats.fuzzy_dups);
2346        assert_eq!(a.stats.hard_negs, b.stats.hard_negs);
2347        assert_eq!(a.stats.uniques, b.stats.uniques);
2348        assert_eq!(a.stats.masters, b.stats.masters);
2349        let _ = std::fs::remove_dir_all(&dir_a);
2350        let _ = std::fs::remove_dir_all(&dir_b);
2351    }
2352
2353    #[test]
2354    fn test_distribute_by_weight_sums_to_total_and_is_proportional() {
2355        // Regression: duplicates used to be sampled only from the last
2356        // batch of masters (`last_batch`), so for any entity spanning 2+
2357        // batches, earlier batches never received a single duplicate —
2358        // the observed symptom was the overall duplicate rate silently
2359        // drifting down as dataset size grew past `BATCH_SIZE`. Fixed by
2360        // spreading each noise_type's count across every batch in
2361        // proportion to its size via this helper.
2362        let weights = [500_000usize, 216_666];
2363        let counts = distribute_by_weight(283_334, &weights);
2364        assert_eq!(counts.len(), 2);
2365        assert_eq!(counts.iter().sum::<usize>(), 283_334);
2366        // Every batch gets a non-zero share proportional to its size —
2367        // in particular the *first* batch (which the old code starved
2368        // entirely) must receive the majority share here.
2369        assert!(counts[0] > 0, "first batch got zero duplicates");
2370        assert!(counts[1] > 0, "last batch got zero duplicates");
2371        assert!(counts[0] > counts[1]);
2372
2373        // Edge cases: zero total, zero weights.
2374        assert_eq!(distribute_by_weight(0, &weights), vec![0, 0]);
2375        assert_eq!(distribute_by_weight(100, &[0, 0]), vec![0, 0]);
2376    }
2377
2378    #[test]
2379    fn test_entity_prefix_unique_per_index() {
2380        let a = entity_prefix(0);
2381        let b = entity_prefix(1);
2382        assert_ne!(a, b);
2383        assert_eq!(a, "E00000");
2384        assert_eq!(entity_prefix(42), "E00042");
2385    }
2386
2387    #[test]
2388    fn test_run_pipeline_hell_multi_pass_touches_multiple_columns() {
2389        // At `hell` (3 noise passes, see `schema::DifficultySettings::
2390        // passes`), at least one duplicate copy should differ from its
2391        // master on 2+ distinct columns — a single pass only ever touches
2392        // the columns its one chosen noise_type's predicate matches
2393        // (`noise_type_targets_column`), so seeing 2+ changed columns on
2394        // the same row is direct evidence that multiple independent passes
2395        // actually landed on it, not just a coincidence of one pass hitting
2396        // a multi-column noise_type like "typo".
2397        let dir = scratch_dir("hell_multipass");
2398        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2399        let mut ctx = Context::new(
2400            "kyc",
2401            "en",
2402            &manifest_dir().join("assets/pools").to_string_lossy(),
2403        )
2404        .expect("load context");
2405        let config = build_pipeline_config(
2406            "kyc",
2407            5000,
2408            7,
2409            "hell",
2410            0.1,
2411            0.3,
2412            &schema,
2413            "pipeline_test_hell_multipass",
2414            "ipc",
2415            false,
2416            "ipc",
2417            None,
2418        )
2419        .expect("build config");
2420        ctx.enable_watermark(&config.domain, config.size, config.seed);
2421        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2422
2423        let read_all = |path: &str| -> RecordBatch {
2424            let file = std::fs::File::open(path).unwrap();
2425            let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
2426            let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2427            let schema = batches[0].schema();
2428            let refs: Vec<Vec<&dyn Array>> = (0..schema.fields().len())
2429                .map(|i| batches.iter().map(|b| b.column(i).as_ref()).collect())
2430                .collect();
2431            let cols: Vec<ArrayRef> = refs
2432                .iter()
2433                .map(|r| arrow::compute::concat(r).unwrap())
2434                .collect();
2435            RecordBatch::try_new(schema, cols).unwrap()
2436        };
2437        let ds = read_all(&output.output_files[0]);
2438        let gt = read_all(&output.gt_file);
2439
2440        let ds_rid = ds.column_by_name("record_id").unwrap().as_string::<i32>();
2441        let ds_mid = ds.column_by_name("master_id").unwrap().as_string::<i32>();
2442        let gt_rid = gt.column_by_name("record_id").unwrap().as_string::<i32>();
2443        let gt_mt = gt.column_by_name("match_type").unwrap().as_string::<i32>();
2444
2445        let mut match_type_by_rid: HashMap<&str, &str> = HashMap::new();
2446        for i in 0..gt.num_rows() {
2447            match_type_by_rid.insert(gt_rid.value(i), gt_mt.value(i));
2448        }
2449
2450        // The true master row for a given master_id is whichever row was
2451        // written *first* (base rows are always written before any of their
2452        // duplicate copies, within a batch) — not "whichever row happens to
2453        // be labeled `exact_dup`": a cluster with only fuzzy copies has no
2454        // `exact_dup` row at all (a base row only reads `exact_dup` when the
2455        // cluster actually contains a byte-identical copy — see `gt.rs`).
2456        let mut master_row_by_mid: HashMap<&str, usize> = HashMap::new();
2457        for i in 0..ds.num_rows() {
2458            master_row_by_mid.entry(ds_mid.value(i)).or_insert(i);
2459        }
2460
2461        let compare_cols: Vec<usize> = (0..ds.num_columns())
2462            .filter(|&i| {
2463                let field = ds.schema().field(i).clone();
2464                matches!(field.data_type(), DataType::Utf8)
2465                    && !matches!(
2466                        field.name().as_str(),
2467                        "record_id" | "domain" | "entity_type" | "master_id"
2468                    )
2469                    && !field.name().ends_with("_id")
2470            })
2471            .collect();
2472
2473        let mut max_changed_cols = 0usize;
2474        for i in 0..ds.num_rows() {
2475            if match_type_by_rid.get(ds_rid.value(i)) != Some(&"fuzzy_dup") {
2476                continue;
2477            }
2478            let Some(&master_i) = master_row_by_mid.get(ds_mid.value(i)) else {
2479                continue;
2480            };
2481            let mut changed = 0usize;
2482            for &col_i in &compare_cols {
2483                let col = ds.column(col_i).as_string::<i32>();
2484                let a = (!col.is_null(i)).then(|| col.value(i));
2485                let b = (!col.is_null(master_i)).then(|| col.value(master_i));
2486                if a != b {
2487                    changed += 1;
2488                }
2489            }
2490            max_changed_cols = max_changed_cols.max(changed);
2491        }
2492
2493        assert!(
2494            max_changed_cols >= 2,
2495            "expected at least one fuzzy_dup row to differ from its master on \
2496             2+ columns (evidence of multiple noise passes), got max {max_changed_cols}"
2497        );
2498
2499        let _ = std::fs::remove_dir_all(&dir);
2500    }
2501
2502    #[test]
2503    fn test_record_id_and_pad_string_roundtrip() {
2504        let rid = record_id_string(7);
2505        assert_eq!(rid, "R-0000000000007");
2506        let pad = pad_string(7);
2507        assert_eq!(pad, "0000000000007");
2508        let strs = record_id_strs(0..3);
2509        assert_eq!(strs.len(), 3);
2510        assert_eq!(strs[0], record_id_string(0));
2511        assert_eq!(strs[2], record_id_string(2));
2512    }
2513
2514    #[test]
2515    fn test_parse_record_idx_roundtrip() {
2516        for i in [0usize, 7, 1_234_567, 300_000_000] {
2517            assert_eq!(parse_record_idx(&record_id_string(i)), Some(i as u64));
2518        }
2519        assert_eq!(parse_record_idx("garbage"), None);
2520        assert_eq!(parse_record_idx("R-1"), None);
2521        assert_eq!(parse_record_idx("X-0000000000007"), None);
2522    }
2523
2524    #[test]
2525    fn test_pack_master_key_roundtrip_and_ordering() {
2526        let mid = |e: usize, i: usize| format!("{}-{}", entity_prefix(e), pad_string(i));
2527        assert!(pack_master_key(&mid(0, 0)).is_some());
2528        // Same entity, ascending local index -> ascending packed key.
2529        assert!(pack_master_key(&mid(0, 1)).unwrap() > pack_master_key(&mid(0, 0)).unwrap());
2530        // Different entity always outranks any local index within a lower
2531        // entity, exactly like comparing the original fixed-width strings.
2532        assert!(
2533            pack_master_key(&mid(1, 0)).unwrap()
2534                > pack_master_key(&mid(0, 9_999_999_999_999)).unwrap()
2535        );
2536        assert_eq!(pack_master_key("garbage"), None);
2537        assert_eq!(pack_master_key("HN-0000000000007"), None);
2538        assert_eq!(pack_master_key("CANARY-000-0-abc"), None);
2539    }
2540
2541    // ── perf-hunt H1 isolated measurement (hunt2407.md) ─────────────────
2542    //
2543    // Not a correctness test -- #[ignore]'d so it never runs in normal
2544    // `cargo test`. Measures whether `apply_noise_with_retry`'s retry-zip
2545    // loop and `apply_extra_pass_per_row`'s pass-zip loop (both iterate
2546    // `0..num_columns()`, not just the noise_type's actual target_cols) cost
2547    // scales with TOTAL column count C, holding the number of noised
2548    // columns T=1 constant -- H1's hypothesis. Run with:
2549    //   cargo test --release pipeline::tests::hunt_h1_zip_cost_vs_column_count -- --ignored --nocapture
2550    #[test]
2551    #[ignore]
2552    fn hunt_h1_zip_cost_vs_column_count() {
2553        use arrow::array::StringArray;
2554        use arrow::datatypes::{DataType, Field, Schema};
2555
2556        fn make_batch(n_rows: usize, n_cols: usize) -> RecordBatch {
2557            let fields: Vec<Field> = (0..n_cols)
2558                .map(|i| Field::new(format!("col{i}"), DataType::Utf8, false))
2559                .collect();
2560            let cols: Vec<ArrayRef> = (0..n_cols)
2561                .map(|_| {
2562                    Arc::new(StringArray::from(
2563                        (0..n_rows)
2564                            .map(|_| "abcdefghijklmnop".to_string())
2565                            .collect::<Vec<_>>(),
2566                    )) as ArrayRef
2567                })
2568                .collect();
2569            RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
2570        }
2571
2572        const N_ROWS: usize = 100_000;
2573        const N_REPS: usize = 5;
2574        // 3 noise_passes total (hell tier): 1 first pass + 2 extra passes,
2575        // matching `apply_noise_with_retry` + `apply_extra_pass_per_row`
2576        // call pattern in `run_pipeline`'s dup-generation closure.
2577        const EXTRA_PASSES: usize = 2;
2578
2579        eprintln!("\n[hunt_h1] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
2580        eprintln!("[hunt_h1] C=total columns, T=1 noised column (fixed) -- H1 predicts cost ∝ C");
2581
2582        for &n_cols in &[8usize, 15, 23] {
2583            let plan_cols = vec!["col0".to_string()];
2584            let all_types: Vec<&str> = vec!["typo"];
2585            let mut total = std::time::Duration::ZERO;
2586
2587            for rep in 0..N_REPS {
2588                let batch = make_batch(N_ROWS, n_cols);
2589                let mut rng = Rng::new(42 + rep as u64);
2590                let t0 = std::time::Instant::now();
2591
2592                let (noisy, mut unchanged, _) =
2593                    apply_noise_with_retry(&batch, "typo", &plan_cols, &mut rng, &[]).unwrap();
2594                let mut current = noisy;
2595                for _ in 0..EXTRA_PASSES {
2596                    current = apply_extra_pass_per_row(
2597                        &current,
2598                        &mut unchanged,
2599                        &all_types,
2600                        &mut rng,
2601                        &[],
2602                    )
2603                    .unwrap();
2604                }
2605                std::hint::black_box(&current);
2606                total += t0.elapsed();
2607            }
2608
2609            let avg = total / N_REPS as u32;
2610            eprintln!(
2611                "[hunt_h1] C={n_cols:2}  avg={:>8.3}ms  per_column={:>7.4}ms  per_row_per_column={:>10.6}us",
2612                avg.as_secs_f64() * 1000.0,
2613                avg.as_secs_f64() * 1000.0 / n_cols as f64,
2614                avg.as_secs_f64() * 1_000_000.0 / (N_ROWS * n_cols) as f64,
2615            );
2616        }
2617    }
2618
2619    // ── perf-hunt kyc residual factor (hunt2407.md, corrected) ──────────
2620    //
2621    // The first `hunt_h1_*` measurement above got two things wrong when
2622    // trying to explain kyc/hell's real-world ~4-5x slowdown vs nonprofit:
2623    // (1) it assumed `noise_passes = 3` at `hell` -- the real value is
2624    // `DifficultySettings::passes = 8` (schema.rs); (2) it used generic
2625    // `col0..colN` column names with `plan_cols=[]` for the "extra pass"
2626    // calls, but `noise_type_targets_column` is name-based -- generic names
2627    // match none of hell's 9 categories, so those extra passes were
2628    // silently no-ops regardless of `EXTRA_PASSES`. This measurement uses
2629    // the REAL kyc.natural_person (23 cols) and nonprofit.donation (7 cols,
2630    // the weight-dominant nonprofit entity, weight=200) column NAMES loaded
2631    // from `schemas/*.json`, hell's real 9-category `noise_types` list, and
2632    // the real pass count (8) -- reproducing genuine per-category column
2633    // matching via `noise_type_targets_column`. Run with:
2634    //   cargo test --release pipeline::tests::hunt_h1_residual_kyc_vs_nonprofit -- --ignored --nocapture
2635    #[test]
2636    #[ignore]
2637    fn hunt_h1_residual_kyc_vs_nonprofit() {
2638        use arrow::array::StringArray;
2639        use arrow::datatypes::{DataType, Field, Schema};
2640
2641        fn make_batch(n_rows: usize, col_names: &[&str]) -> RecordBatch {
2642            let fields: Vec<Field> = col_names
2643                .iter()
2644                .map(|n| Field::new(*n, DataType::Utf8, true))
2645                .collect();
2646            let cols: Vec<ArrayRef> = col_names
2647                .iter()
2648                .map(|_| {
2649                    Arc::new(StringArray::from(
2650                        (0..n_rows)
2651                            .map(|_| "abcdefghijklmnop testword".to_string())
2652                            .collect::<Vec<_>>(),
2653                    )) as ArrayRef
2654                })
2655                .collect();
2656            RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
2657        }
2658
2659        const N_ROWS: usize = 100_000;
2660        const N_REPS: usize = 5;
2661        const HELL_TYPES: &[&str] = &[
2662            "typo_aggressive",
2663            "visual",
2664            "unicode_pollution",
2665            "names",
2666            "dates",
2667            "identifiers",
2668            "addresses",
2669            "companies",
2670            "extra",
2671        ];
2672        const HELL_PASSES: usize = 8; // schema.rs DifficultySettings::passes for "hell"
2673
2674        let kyc_cols: &[&str] = &[
2675            "party_id",
2676            "given_name",
2677            "family_name",
2678            "middle_name",
2679            "birth_date",
2680            "birth_place",
2681            "nationality",
2682            "email_address",
2683            "mobile_phone",
2684            "personal_administrative_number",
2685            "tax_id",
2686            "residential_address_line1",
2687            "residential_city",
2688            "residential_postal_code",
2689            "residential_state",
2690            "residential_country",
2691            "occupation",
2692            "risk_score",
2693            "document_type",
2694            "document_number",
2695            "source_system",
2696        ];
2697        let nonprofit_cols: &[&str] = &[
2698            "donation_id",
2699            "donor_id",
2700            "campaign_id",
2701            "donation_date",
2702            "payment_method",
2703            "source_system",
2704        ];
2705
2706        eprintln!("\n[hunt_h1_residual] N_ROWS={N_ROWS} N_REPS={N_REPS} HELL_PASSES={HELL_PASSES}");
2707
2708        // Per-category breakdown for kyc only (the slow one) -- one rep,
2709        // timing each of the 9 categories' apply_noise_with_retry call
2710        // individually (fresh copy of the batch each time, not chained) to
2711        // find which category(ies) dominate the 18s total.
2712        {
2713            let batch = make_batch(N_ROWS, kyc_cols);
2714            eprintln!("[hunt_h1_residual] per-category breakdown on kyc.natural_person:");
2715            for &ttype in HELL_TYPES {
2716                let mut rng = Rng::new(42);
2717                let t0 = std::time::Instant::now();
2718                let (_, _, tcols) =
2719                    apply_noise_with_retry(&batch, ttype, &[], &mut rng, &[]).unwrap();
2720                eprintln!(
2721                    "[hunt_h1_residual]   {ttype:20} {:>8.3}ms  target_cols={:?}",
2722                    t0.elapsed().as_secs_f64() * 1000.0,
2723                    tcols,
2724                );
2725            }
2726        }
2727
2728        for (label, col_names) in [
2729            ("kyc.natural_person", kyc_cols),
2730            ("nonprofit.donation", nonprofit_cols),
2731        ] {
2732            let mut total = std::time::Duration::ZERO;
2733            let mut total_target_hits = 0usize;
2734
2735            for rep in 0..N_REPS {
2736                let batch = make_batch(N_ROWS, col_names);
2737                let mut rng = Rng::new(42 + rep as u64);
2738                let t0 = std::time::Instant::now();
2739
2740                // First pass: same as real generation, one explicit
2741                // noise_type (heuristic-matched via plan_cols=[]).
2742                let (noisy, mut unchanged, t0cols) =
2743                    apply_noise_with_retry(&batch, "identifiers", &[], &mut rng, &[]).unwrap();
2744                total_target_hits += t0cols.len();
2745                let mut current = noisy;
2746                for _ in 1..HELL_PASSES {
2747                    current = apply_extra_pass_per_row(
2748                        &current,
2749                        &mut unchanged,
2750                        HELL_TYPES,
2751                        &mut rng,
2752                        &[],
2753                    )
2754                    .unwrap();
2755                }
2756                std::hint::black_box(&current);
2757                total += t0.elapsed();
2758            }
2759
2760            let avg = total / N_REPS as u32;
2761            eprintln!(
2762                "[hunt_h1_residual] {label:20} C={:2}  avg={:>8.3}ms  avg_first_pass_target_cols={:.1}",
2763                col_names.len(),
2764                avg.as_secs_f64() * 1000.0,
2765                total_target_hits as f64 / N_REPS as f64,
2766            );
2767        }
2768    }
2769}