Skip to main content

dupehell_core/
pipeline.rs

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