dupehell_core/

pipeline.rs

// DupeHell -- MIT License
//
// Synthetic multi-domain dataset generator for record linkage benchmarking.
// No liability for misuse.

use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::{
    Array, ArrayRef, AsArray, StringArray, StringBuilder, UInt64Array, UInt64Builder,
};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use serde::Deserialize;

use crate::context::Context;
use crate::rng::Rng;

// ── Config types ───────────────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct PipelineConfig {
    pub domain: String,
    pub size: usize,
    pub seed: u64,
    pub difficulty: String,
    pub output_format: String,
    pub run_id: String,

    pub entity_plans: Vec<EntityPlan>,
    pub hard_neg_types: Vec<HnTypeConfig>,

    pub hard_neg_ratio: f64,

    /// Number of independent noise passes applied to each duplicate copy —
    /// see the doc comment on `schema::DifficultySettings::passes` for why
    /// this (not the noise_types category count) is the actual difficulty
    /// lever between tiers.
    pub noise_passes: usize,

    pub graph_enabled: bool,
    pub graph_format: String,

    /// `--skip-ground-truth`: skip ground-truth classification entirely and
    /// don't write the `_ground_truth` file at all (see
    /// `gt::GtAccumulator`). The GT file is about a third of the dataset's
    /// size, so skipping it saves that much write I/O plus the per-row
    /// classification — for stress-test workloads that only need the
    /// dataset itself. Defaults to `false` (absent from
    /// the JSON `build_pipeline_config` constructs) so every existing
    /// caller keeps producing ground truth unchanged. Incompatible with
    /// `--graph`, which needs the post-GT cluster map to emit duplicate-
    /// cluster edges — validated in `main.rs`, not here.
    #[serde(default)]
    pub skip_ground_truth: bool,
}

#[derive(Debug, Deserialize)]
pub struct EntityPlan {
    pub name: String,
    pub n_base: usize,
    pub identifier_col: Option<String>,
    pub columns_json: String,
    pub noise_types: Vec<NoisePlanEntry>,
    pub fk_remaps: Vec<FkRemapEntry>,
    /// Set by `schema::build_pipeline_config` for the FK-target entities
    /// added when `--only-entity` is active (see `FK_POOL_CAP`): this
    /// entity's rows are generated only far enough to build its identifier
    /// pool for other entities' FK remaps, never noised, written to the
    /// dataset/graph output, or fed to ground truth. Always `false` on the
    /// normal (no `--only-entity`) path.
    #[serde(default)]
    pub pool_only: bool,
}

/// Cap on how many identifiers of an entity are retained as a pool for other
/// entities' FK remaps to sample from (`fk_remap.rs` samples uniformly via
/// `take`, so it never needs to see every identifier the entity produced).
/// Also used by `schema::build_pipeline_config` to size the pool-only
/// entities it synthesizes for `--only-entity` (no need to generate more
/// rows of an FK target than this pool will ever draw from).
pub(crate) const FK_POOL_CAP: usize = 200_000;

#[derive(Debug, Deserialize)]
pub struct NoisePlanEntry {
    pub noise_type: String,
    pub columns: Vec<String>,
    pub count: usize,
}

#[derive(Debug, Deserialize)]
pub struct FkRemapEntry {
    pub source_col: String,
    pub target_entity: String,
    /// Weight written to the `weight` column of `fk`-typed graph edges
    /// produced from this relation (`--graph`). Defaults to 1.0 for schemas
    /// that don't declare it, preserving prior behavior.
    #[serde(default = "default_fk_weight")]
    pub fk_weight: f64,
}

fn default_fk_weight() -> f64 {
    1.0
}

#[derive(Debug, Deserialize)]
pub struct HnTypeConfig {
    pub entity_type: String,
    pub config_json: String,
    pub count: usize,
}

#[derive(Debug)]
pub struct PipelineOutput {
    pub output_files: Vec<String>,
    pub gt_file: String,
    pub stats: PipelineStats,
    pub nodes: Option<String>,
    pub edges: Option<String>,
}

#[derive(Debug)]
pub struct PipelineStats {
    pub total_records: usize,
    /// Cluster members that are genuinely byte-for-byte identical to their
    /// master. See `PipelineStats::fuzzy_dups` for the noised counterpart.
    pub exact_dups: usize,
    /// Cluster members that are duplicates of a master but differ from it on
    /// at least one column (the assigned noise actually changed something).
    pub fuzzy_dups: usize,
    pub hard_negs: usize,
    pub uniques: usize,
    pub masters: usize,
}

/// Carries the running counters that `record_id`/`master_id` generation
/// depend on across a sequence of chunked runs (`--chunk-size`), so chunk
/// N+1 continues exactly where chunk N left off instead of restarting at
/// 0 — genuinely contiguous, global identifiers over the whole logical
/// dataset, with no shard-tag hack needed on top. Defaults (all zero) give
/// the single-run behavior, unchanged.
#[derive(Debug, Clone, Default)]
pub struct ChunkOffsets {
    /// Next `record_id` index to hand out (see `record_id_string`).
    pub rid_offset: usize,
    /// Next local `master_id` index per entity, indexed by the entity's
    /// position in `PipelineConfig::entity_plans` (see `entity_prefix`).
    /// Entities absent from a shorter previous chunk's plan (shouldn't
    /// happen — the same schema/`--only-entity` produces the same plan
    /// shape every chunk) simply default to 0 via `master_base`.
    pub master_local: Vec<usize>,
    /// Next `HN-{:09}` counter value (see `hn_master_id_counter`).
    pub hn_master_counter: u64,
    /// Next `row_offset` per entity (same indexing as `master_local`) for
    /// the generated `_id` columns (`column_gen`'s `PREFIX-0000000042`
    /// fallback): past this entity's base rows *and* its canary rows of
    /// every previous chunk. Without it each chunk numbered its `_id`s from
    /// 0 again, so two unrelated records of different chunks shared one
    /// `party_id`/`customer_id`.
    pub id_local: Vec<usize>,
}

impl ChunkOffsets {
    fn master_base(&self, entity_idx: usize) -> usize {
        self.master_local.get(entity_idx).copied().unwrap_or(0)
    }

    pub(crate) fn id_base(&self, entity_idx: usize) -> usize {
        self.id_local.get(entity_idx).copied().unwrap_or(0)
    }
}

// ── Fixed-width ASCII IDs ────────────────────────────────────────────────

// `record_id`/`pad` are pure functions of the row index — "R-" + 13 digits
// and 13 digits respectively. Earlier this was a struct precomputing both
// for every id up front (bounded per-2-big-allocations, better than one
// `String` per id, but still O(total_ids) resident for the whole run — at
// 200M+ ids that's several GB held for nothing). Computing on demand, only
// for the batch currently being written, removes that fixed cost entirely;
// the transient `Vec<String>` built per batch is bounded by `BATCH_SIZE`
// and freed right after use, same as `batch_mids`/`dup_mids_buf`/`hn_mids`
// elsewhere in this file.
pub(crate) const RID_LEN: usize = 15; // "R-" + 13 digits
pub(crate) const PAD_LEN: usize = 13; // 13 digits

// Every real (non-test) code path now uses `record_id_array`/
// `master_id_array_*` or, on scalar (one-id-at-a-time) call
// sites, `append_record_id` — a buffer-reusing sibling
// with the same output, no per-id `String` allocation. `record_id_string`
// itself survives only as a `#[cfg(test)]` fixture helper now — kept
// `#[cfg(test)]` rather than deleted since test fixtures across
// `graph_gen.rs`/`gt.rs` still build expected ids with it.
#[cfg(test)]
#[inline]
pub(crate) fn record_id_string(i: usize) -> String {
    let mut buf = [0u8; RID_LEN];
    buf[0] = b'R';
    buf[1] = b'-';
    let mut n = i;
    for j in (2..RID_LEN).rev() {
        buf[j] = b'0' + (n % 10) as u8;
        n /= 10;
    }
    String::from_utf8(buf.to_vec()).expect("record_id_string: buffer is 'R', '-' and ASCII digits")
}
#[cfg(test)]
#[inline]
pub(crate) fn pad_string(i: usize) -> String {
    let mut buf = [0u8; PAD_LEN];
    let mut n = i;
    for j in (0..PAD_LEN).rev() {
        buf[j] = b'0' + (n % 10) as u8;
        n /= 10;
    }
    String::from_utf8(buf.to_vec()).expect("pad_string: buffer is always ASCII digits")
}

#[cfg(test)]
pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
    range.map(record_id_string).collect()
}

/// Writes `val` right-aligned, zero-padded to `width` ASCII digits, at the
/// end of `buf` (grows `buf` by exactly `width` bytes). Shared digit-writing
/// core for `record_id_array`/`master_id_array_*`/`hn_master_id_array` —
/// same arithmetic as `record_id_string`/`pad_string`, just writing into a
/// caller-owned buffer instead of allocating its own.
#[inline]
pub(crate) fn write_digits(buf: &mut Vec<u8>, val: u64, width: usize) {
    let start = buf.len();
    buf.resize(start + width, b'0');
    let mut n = val;
    for j in (0..width).rev() {
        buf[start + j] = b'0' + (n % 10) as u8;
        n /= 10;
    }
}

/// Appends `"R-" + idx` zero-padded to `PAD_LEN` digits onto the end of
/// `buf` — the exact same bytes `record_id_string(idx)` produces, without
/// allocating a `String`. Callers writing many
/// ids in a loop should reuse one `buf` across rows (`buf.clear()` between
/// calls, same convention as `buf_gen::build_string_array`'s row buffer).
#[inline]
pub(crate) fn append_record_id(buf: &mut Vec<u8>, idx: u64) {
    buf.push(b'R');
    buf.push(b'-');
    write_digits(buf, idx, PAD_LEN);
}

/// Builds a `record_id` `StringArray` for `range` directly (one shared
/// buffer via `buf_gen::build_string_array`, no per-row `String`) instead of
/// collecting `Vec<String>` and recopying through `StringArray::
/// from_iter_values` — same bytes, same order.
pub(crate) fn record_id_array(range: std::ops::Range<usize>) -> ArrayRef {
    let n = range.len();
    let mut i = range.start as u64;
    crate::buf_gen::build_string_array(n, RID_LEN, |buf| {
        buf.push(b'R');
        buf.push(b'-');
        write_digits(buf, i, PAD_LEN);
        i += 1;
    })
}

/// Builds a `master_id` `StringArray` for the contiguous global-index range
/// `start..start+n` (base entity batches: every row in the batch gets the
/// next master index in order).
fn master_id_array_range(prefix: &str, start: usize, n: usize) -> ArrayRef {
    let width = prefix.len() + 1 + PAD_LEN;
    let mut i = start as u64;
    crate::buf_gen::build_string_array(n, width, |buf| {
        buf.extend_from_slice(prefix.as_bytes());
        buf.push(b'-');
        write_digits(buf, i, PAD_LEN);
        i += 1;
    })
}

/// Builds a `master_id` `StringArray` from an arbitrary (non-contiguous)
/// sequence of global indices — duplicate-copy rows, whose master is
/// whichever base row `indices` sampled, not a running counter. `indices`
/// must yield exactly `n` values.
fn master_id_array_from_indices(
    prefix: &str,
    mut indices: impl Iterator<Item = usize>,
    n: usize,
) -> ArrayRef {
    let width = prefix.len() + 1 + PAD_LEN;
    crate::buf_gen::build_string_array(n, width, |buf| {
        let idx = indices
            .next()
            .expect("master_id_array_from_indices: fewer than n indices");
        buf.extend_from_slice(prefix.as_bytes());
        buf.push(b'-');
        write_digits(buf, idx as u64, PAD_LEN);
    })
}

/// Builds an `"HN-{:09}"` `master_id` `StringArray` for `n` rows, counting
/// up from `start` — mirrors the previous per-row `format!("HN-{:09}", id)`
/// loop, same values, same order.
fn hn_master_id_array(start: u64, n: usize) -> ArrayRef {
    let mut id = start;
    crate::buf_gen::build_string_array(n, 3 + HN_ID_DIGITS, |buf| {
        buf.extend_from_slice(b"HN-");
        write_digits(buf, id, HN_ID_DIGITS);
        id += 1;
    })
}

/// Inverse of `record_id_string`: parses `"R-0000000000007"` back to `7`.
/// Returns `None` if `rid` isn't exactly `RID_LEN` bytes starting with
/// `"R-"` followed by `PAD_LEN` ASCII digits (never expected in practice —
/// every `record_id` in a duplicated cluster is produced by
/// `record_id_string`/`record_id_strs`, never hand-built).
pub(crate) fn parse_record_idx(rid: &str) -> Option<u64> {
    if rid.len() != RID_LEN || !rid.starts_with("R-") {
        return None;
    }
    rid[2..].parse().ok()
}

/// Packs a `master_id` of the fixed shape `"{entity_prefix}-{pad_string}"`
/// (e.g. `"E00042-0000000001234"`) into a single `u64`: `entity_idx` in the
/// high bits, `local_idx` in the low 44 bits. `entity_idx` is at most 5
/// digits (`entity_prefix`'s format), so it fits comfortably above bit 44;
/// `local_idx` is at most `PAD_LEN` (13) digits, which fits in 44 bits
/// (2^44 ≈ 1.76e13). Preserves the same ordering as a lexicographic
/// comparison of the original fixed-width string, since both segments are
/// zero-padded to a constant width — callers that need the old
/// `master_ids.sort()` determinism can just sort/compare the packed `u64`.
/// Returns `None` if `mid` isn't `"E" + 5 digits + "-" + PAD_LEN digits`
/// (never expected for an entity master_id — HN-/CANARY- ids are the only
/// other shapes, and they're never part of a duplicate cluster).
pub(crate) fn pack_master_key(mid: &str) -> Option<u64> {
    let bytes = mid.as_bytes();
    if bytes.len() != 1 + 5 + 1 + PAD_LEN || bytes[0] != b'E' || bytes[6] != b'-' {
        return None;
    }
    let entity_idx: u64 = mid[1..6].parse().ok()?;
    let local_idx: u64 = mid[7..7 + PAD_LEN].parse().ok()?;
    Some(pack_master_parts(entity_idx, local_idx))
}

/// Same packing as [`pack_master_key`], straight from the numbers the
/// pipeline formats into the master_id (`entity_prefix(entity_idx)` +
/// `pad_string(local_idx)`) — no string parsing.
pub(crate) fn pack_master_parts(entity_idx: u64, local_idx: u64) -> u64 {
    (entity_idx << 44) | local_idx
}

/// Splits `total` proportionally across `weights` (largest-remainder
/// method), so the sum of the result is always exactly `total`.
///
/// Used to spread each noise_type's planned duplicate count across every
/// batch of masters in proportion to batch size, instead of assigning it
/// all to a single batch — see the doc comment on the "Dups" section in
/// `run_pipeline_with_progress` for why that matters.
fn distribute_by_weight(total: usize, weights: &[usize]) -> Vec<usize> {
    let weight_sum: usize = weights.iter().sum();
    if weight_sum == 0 || total == 0 {
        return vec![0; weights.len()];
    }
    let raw: Vec<f64> = weights
        .iter()
        .map(|&w| total as f64 * w as f64 / weight_sum as f64)
        .collect();
    let mut counts: Vec<usize> = raw.iter().map(|&r| r as usize).collect();
    let assigned: usize = counts.iter().sum();
    let mut remainder = total - assigned;
    let mut order: Vec<usize> = (0..counts.len()).collect();
    order.sort_by(|&a, &b| {
        let fa = raw[a] - counts[a] as f64;
        let fb = raw[b] - counts[b] as f64;
        fb.partial_cmp(&fa).unwrap_or(std::cmp::Ordering::Equal)
    });
    for &i in &order {
        if remainder == 0 {
            break;
        }
        counts[i] += 1;
        remainder -= 1;
    }
    counts
}

// ── Entity prefix ──────────────────────────────────────────────────────────

/// Per-entity master_id prefix, derived from the entity's position in
/// `config.entity_plans` rather than a hash of its name: a name hash mod
/// 100000 had a real (if small) chance of two entities in the same domain
/// colliding on the same prefix, which would make `gt.rs`'s master_id
/// counting conflate masters from different entity types. An index is
/// unique by construction — zero collision risk regardless of domain size.
pub(crate) fn entity_prefix(entity_idx: usize) -> String {
    format!("E{:05}", entity_idx)
}

// ── FK / HN pool types ─────────────────────────────────────────────────────

type FkPoolMap = HashMap<String, RecordBatch>;

struct HnPool {
    batch: RecordBatch,
    total_count: usize,
    // Global record_ids aligned positionally with `batch` rows (only populated
    // when graph output is enabled). Used to emit `hard_neg` edges.
    record_ids: ArrayRef,
}

// ── Noise column matching ─────────────────────────────────────────────────

/// Column-name fragments for columns holding a person's name as free text,
/// even when the column isn't literally called `*_name` (e.g. `operator`,
/// `technician`, populated from the `first_name` pool across several
/// domain schemas).
const PERSON_NAME_WORDS: &[&str] = &[
    "name",
    // Not bare "first"/"last": those also matched non-name columns like
    // `first_seen`/`last_active`/`last_login`/`last_updated`, sending them
    // through char-level typo/homoglyph noise and injecting letters into
    // what are actually timestamps. Every legitimate name
    // column across the 40 domain schemas is literally `*first_name*` /
    // `*last_name*` (e.g. `agent_first_name`, `pat_first_name`), so this
    // loses no real coverage.
    "first_name",
    "last_name",
    "given",
    "family",
    "middle",
    "operator",
    "technician",
    "inspector",
    "claimant",
    "consignee",
    "investigator",
    "collected_by",
    "performed_by",
    "assigned_to",
];

/// Column-name fragments for columns holding a company/organization name,
/// even when the column isn't literally called `*_name`/`company*` (e.g.
/// `supplier`, `manufacturer`, populated from the `company` pool across
/// several domain schemas).
const COMPANY_NAME_WORDS: &[&str] = &[
    "company",
    "legal",
    "trading",
    "name",
    "supplier",
    "manufacturer",
    "sponsor",
    "institution",
    "journal",
    "funder",
    "employer",
    "law_firm",
    "affiliation",
    "reported_by",
    "buyer",
    "shipper",
    "airline",
];

/// Does `lower` match one of `PERSON_NAME_WORDS`' fragments *other than* the
/// bare `"name"` catch-all — i.e. is it identifiable as a person-name column
/// via a specific signal (`given`, `family`, `first_name`, `inspector`...),
/// not just because its name happens to end in "name"?
///
/// Used to keep the `"name"` fragment both lists share (needed to catch
/// columns like `judge_name`/`director_name` with no more specific word)
/// from making `given_name`/`company_name` cross-match each other's noise
/// category — see `is_specifically_company_name` (kyc's `given_name`/`family_name`/
/// `middle_name` were getting `companies` noise — `drop_legal_form` etc —
/// applied to them, purely because `COMPANY_NAME_WORDS` also contains
/// `"name"`).
fn is_specifically_person_name(lower: &str) -> bool {
    PERSON_NAME_WORDS
        .iter()
        .any(|&w| w != "name" && lower.contains(w))
}

/// Company-side mirror of `is_specifically_person_name`: matches a
/// `COMPANY_NAME_WORDS` fragment other than the bare `"name"` catch-all.
fn is_specifically_company_name(lower: &str) -> bool {
    COMPANY_NAME_WORDS
        .iter()
        .any(|&w| w != "name" && lower.contains(w))
}

/// Does this noise category ever target a column with this (lowercased) name?
///
/// Pure name-pattern predicate shared with `estimate_difficulty` (see
/// `crate::difficulty`) so the theoretical column-reliability model can never
/// silently drift from what real generation actually does to each column.
pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
    let lower = col_name.to_lowercase();
    match noise_type {
        "typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
        | "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
        | "blocking_fail" => {
            // Skip email-like columns — typo/visual noise destroys '@'.
            // Skip `ip_address` too — it matches the "address" fragment
            // below by name collision, but it isn't a postal address; this
            // char-level noise injected letters into octets.
            if lower.contains("email") || lower.contains("ip_address") {
                return false;
            }
            // Not "phone": that's digit-formatted data, not free text —
            // char-level typo/homoglyph noise injected letters into phone
            // numbers. `identifiers`/`corrupt_phone` already
            // corrupts phone columns digit-aware.
            contains_any(&lower, &["address", "street", "city"])
                || contains_any(&lower, PERSON_NAME_WORDS)
                || contains_any(&lower, COMPANY_NAME_WORDS)
        }
        "names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
            // Exclude columns that are specifically identifiable as a
            // company/organization name (e.g. `company_name`, `legal_name`)
            // — the only reason they'd match PERSON_NAME_WORDS at all is the
            // bare "name" fragment both lists share, not a real person-name
            // signal. See `is_specifically_company_name`.
            contains_any(&lower, PERSON_NAME_WORDS) && !is_specifically_company_name(&lower)
        }
        "dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
            // "time"/"since"/"period" added: the schema's
            // `date`/`datetime` column type already tells the generator a
            // column is a date, but this predicate only ever looked at the
            // *name* -- `departure_time`/`arrival_time` (aviation, travel),
            // `member_since` (hospitality), `pay_period_start`/`_end`
            // (hr), `admission_time`/`discharge_time` (healthcare),
            // `period_start`/`_end` (energy) all carry a real `date`/
            // `datetime` schema type but none contain "date"/"birth"/
            // "incorporation"/"founding" -- silently zero noise coverage
            // from this category regardless of `passes`, on every
            // difficulty tier. Checked against all 40 domains' column
            // names for false-positive collisions before adding (none:
            // the few non-Utf8 hits, e.g. `lifetime_stays`/`lead_time_days`/
            // `on_time_rate`, are ints/floats already excluded upstream by
            // `match_noise_columns`'s Utf8-only filter).
            contains_any(
                &lower,
                &[
                    "date",
                    "birth",
                    "incorporation",
                    "founding",
                    "time",
                    "since",
                    "period",
                ],
            )
        }
        "missing" | "missing_pattern" => {
            contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
        }
        "identifiers"
        | "corrupt_email"
        | "corrupt_phone"
        | "corrupt_national_id"
        | "corrupt_siren"
        | "national_id_corrupt"
        | "phone_corrupt"
        | "email_corrupt"
        | "siren_corrupt" => contains_any(
            &lower,
            &[
                "email",
                "phone",
                "ssn",
                "pan",
                "passport",
                "account",
                "medicare",
                "license",
                "siren",
                "national_id",
                // Added: `difficulty::base_noise_damage` already models a
                // 0.6 damage weight for "registration"-matching columns
                // (e.g. kyc's `registration_number`), but this predicate
                // never actually targeted them -- a real column with zero
                // noise coverage regardless of `passes`, floor-ing hell's
                // f1_max well above its intended ceiling on that schema.
                "registration",
                // "sku"/"barcode"/"imei"/"imsi"/"iccid" added:
                // real sector-specific identifiers this fixed keyword list
                // didn't recognize -- `sku` (ecommerce/fashion/retail/
                // supplychain), `barcode`/`package_barcode` (ecommerce/
                // biotech/logistics), `imei`/`imsi`/`iccid` (telecom
                // `device`, standard cellular identifiers) all had zero
                // noise coverage from this category. Checked against all
                // 40 domains for false-positive collisions before adding
                // (none found).
                "sku",
                "barcode",
                "imei",
                "imsi",
                "iccid",
            ],
        ),
        "extra"
        | "name_null"
        | "dob_null"
        | "blocking_fail_initial"
        | "blocking_fail_partial"
        | "fuzzy_match"
        | "phonetic" => {
            // Same exclusions as the typo/visual arm above:
            // skip email (destroys '@') and `ip_address` (not a postal
            // address despite the name collision), and don't target `phone`
            // with this char-level noise — `identifiers` already handles it
            // digit-aware.
            if lower.contains("email") || lower.contains("ip_address") {
                return false;
            }
            // "bio" added: free-text fields named `bio`/
            // `biography` (social_media, publishing) are exactly the kind
            // of column this category targets, just not covered by the
            // existing "note"/"comment" keywords.
            contains_any(
                &lower,
                &["address", "street", "city", "note", "comment", "bio"],
            ) || contains_any(&lower, PERSON_NAME_WORDS)
                || contains_any(&lower, COMPANY_NAME_WORDS)
        }
        "companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
            // Mirror exclusion: don't apply company-name noise (legal-form
            // dropping, acronym-ing...) to a column that's specifically
            // identifiable as a person's name (e.g. `given_name`,
            // `family_name`) just because it also contains "name". See
            // `is_specifically_person_name`.
            contains_any(&lower, COMPANY_NAME_WORDS) && !is_specifically_person_name(&lower)
        }
        "addresses" | "address_scramble" | "postal_corrupt" => {
            // `ip_address` matches "address" by name collision but isn't a
            // postal address — scrambling/postal-corrupting it doesn't make
            // sense (see the same exclusion on the typo/visual arm
            // above). `email_address` is the same collision: it
            // contains "address" but is an email, not a postal address —
            // `identifiers`/`corrupt_email` already handles it (confirmed
            // empirically: kyc's `email_address` was
            // getting postal `address_scramble`/`postal_corrupt` noise).
            //
            // "state"/"country" added: `difficulty::base_noise_damage`
            // already models a 0.5 damage weight for them (same bucket as
            // address/street/city/postal), but this predicate never
            // actually targeted them -- e.g. kyc's `residential_state` /
            // `residential_country` / `registered_country` had zero real
            // noise coverage regardless of `passes`.
            // "location" added: `social_media::profile`'s
            // `location` field is semantically an address/place field but
            // wasn't covered by any existing keyword. Checked against all
            // 40 domains for false-positive substring collisions before
            // adding (none — the only "location"-containing columns are
            // genuinely location-like: `work_location`, `branch_location`,
            // `shelf_location`).
            !lower.contains("ip_address")
                && !lower.contains("email")
                && contains_any(
                    &lower,
                    &[
                        "address", "street", "postal", "city", "state", "country", "location",
                    ],
                )
        }
        "exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
        | "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
        | "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
        | "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
        | "combo_ultimate" | "french_address" => {
            // These noise types are handled inline or are no-ops
            false
        }
        // Default: applies to all string columns (except FK columns)
        _ => true,
    }
}

fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
    let mut matched: Vec<String> = Vec::new();
    for field in schema.fields() {
        if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
            continue;
        }
        if exclude_cols.contains(field.name()) {
            continue;
        }
        if noise_type_targets_column(noise_type, field.name()) {
            matched.push(field.name().clone());
        }
    }
    // For swap, keep exactly 2 columns
    if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
        matched.truncate(2);
    }
    matched.sort();
    matched.dedup();
    matched
}

fn contains_any(s: &str, patterns: &[&str]) -> bool {
    patterns.iter().any(|p| s.contains(p))
}

// ── Noise application ──────────────────────────────────────────────────────

/// Returns the noised batch plus the list of columns it actually targeted
/// (resolved from the heuristic when `plan_cols` is empty) — callers use
/// the latter to verify the noise produced a real change, see
/// `unchanged_row_mask`.
fn apply_noise_to_batch(
    rb: &RecordBatch,
    noise_type: &str,
    plan_cols: &[String],
    rng: &mut Rng,
    exclude_cols: &[String],
) -> Result<(RecordBatch, Vec<String>), String> {
    let schema = rb.schema();
    let n_cols = rb.num_columns();

    // Fallback to heuristic when plan_cols is empty
    let target_cols: Vec<String> = if plan_cols.is_empty() {
        match_noise_columns(&schema, noise_type, exclude_cols)
    } else {
        plan_cols.to_vec()
    };

    // Start with None; only clone/modify targeted columns
    let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();

    let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
        new_columns[idx]
            .clone()
            .unwrap_or_else(|| rb.column(idx).clone())
    };

    match noise_type {
        "blocking_fail" => {
            for col_name in &target_cols {
                let lower = col_name.to_lowercase();
                if let Ok(idx) = schema.index_of(col_name) {
                    let col = get_col(&new_columns, idx, rb);
                    if lower.contains("first") || lower.contains("given") {
                        new_columns[idx] =
                            Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
                    } else {
                        new_columns[idx] =
                            Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
                    }
                }
            }
        }
        "swap" | "full_swap" => {
            if target_cols.len() >= 2 {
                let name_a = &target_cols[0];
                let name_b = &target_cols[1];
                if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
                    let a = get_col(&new_columns, idx_a, rb);
                    let b = get_col(&new_columns, idx_b, rb);
                    new_columns[idx_a] = Some(b);
                    new_columns[idx_b] = Some(a);
                }
            }
        }
        "missing_pattern" => {
            let mut phone_idx: Option<usize> = None;
            let mut email_idx: Option<usize> = None;
            for col_name in &target_cols {
                let lower = col_name.to_lowercase();
                if let Ok(idx) = schema.index_of(col_name) {
                    if lower.contains("phone") || lower.contains("mobile") {
                        phone_idx = Some(idx);
                    } else if lower.contains("email") {
                        email_idx = Some(idx);
                    }
                }
            }
            if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
                let phone_arr = get_col(&new_columns, pi, rb);
                let email_arr = get_col(&new_columns, ei, rb);
                let phone_s = phone_arr.as_string::<i32>();
                let email_s = email_arr.as_string::<i32>();
                let n = phone_s.len();
                let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
                for i in 0..n {
                    let pv = if phone_s.is_null(i) {
                        ""
                    } else {
                        phone_s.value(i)
                    };
                    let ev = if email_s.is_null(i) {
                        ""
                    } else {
                        email_s.value(i)
                    };
                    if !ev.is_empty() && rng.next_usize(2) == 0 {
                        builder.append_value("");
                    } else {
                        builder.append_value(pv);
                    }
                }
                new_columns[pi] = Some(Arc::new(builder.finish()));
            }
        }
        _ => {
            for col_name in &target_cols {
                if let Ok(col_idx) = schema.index_of(col_name) {
                    let col = get_col(&new_columns, col_idx, rb);
                    let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
                        .map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
                    new_columns[col_idx] = Some(result);
                }
            }
        }
    }

    // Fill remaining from original, unwrap all
    let final_columns: Vec<ArrayRef> = (0..n_cols)
        .map(|i| {
            new_columns[i]
                .clone()
                .unwrap_or_else(|| rb.column(i).clone())
        })
        .collect();

    let batch =
        RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))?;
    Ok((batch, target_cols))
}

/// Per-row mask (`true` = still unchanged): compares `orig` and `noised` on
/// every column in `target_cols`, treating null == null as equal. A `true`
/// row means the noise application was a complete no-op for it — either
/// because its target columns were already null (a "corrupted" null stays
/// null) or because the noise function's own internal randomness happened
/// not to produce a visible variant (several functions in `noise/names.rs`
/// and `noise/extra.rs` only mutate a row some fraction of the time).
/// Consumed by the retry loop in the dup-generation closure to guarantee a
/// duplicate copy that was assigned a noise type actually differs from its
/// master, instead of silently staying a byte-for-byte duplicate.
fn unchanged_row_mask(
    orig: &RecordBatch,
    noised: &RecordBatch,
    target_cols: &[String],
) -> Result<Vec<bool>, String> {
    let n = orig.num_rows();
    unchanged_row_mask_within(orig, noised, target_cols, &vec![true; n])
}

/// Same as [`unchanged_row_mask`], but only re-tests rows where `only[i]`
/// is `true` — every other row keeps its input value unconditionally.
/// Used by `apply_noise_with_retry`'s retry loop:
/// after a retry, `zip(mask, retried, noisy)` only ever overwrites rows
/// where `unchanged[i]` was `true` going in, so for any row where it was
/// `false`, `noisy` coming out of that `zip` is byte-for-byte the same
/// `noisy` that went in on every target column — re-comparing it to `orig`
/// is guaranteed to reproduce the same `false` verdict. The previous
/// implementation re-scanned all `n` rows on every attempt regardless,
/// including rows already proven changed on attempt 1.
fn unchanged_row_mask_within(
    orig: &RecordBatch,
    noised: &RecordBatch,
    target_cols: &[String],
    only: &[bool],
) -> Result<Vec<bool>, String> {
    let n = orig.num_rows();
    // Rows outside `only` are not re-tested: they keep `false` (already
    // known changed), matching the invariant documented above.
    let mut changed: Vec<bool> = only.iter().map(|&keep| !keep).collect();
    let schema = orig.schema();
    for col_name in target_cols {
        let Ok(idx) = schema.index_of(col_name) else {
            continue;
        };
        let a = orig.column(idx).as_string::<i32>();
        let b = noised.column(idx).as_string::<i32>();
        // `i` indexes `a`/`b` (Arrow arrays, not iterable by `.iter_mut()`
        // alongside `changed`) and `changed` in lockstep.
        #[allow(clippy::needless_range_loop)]
        for i in 0..n {
            if changed[i] {
                continue;
            }
            let av = (!a.is_null(i)).then(|| a.value(i));
            let bv = (!b.is_null(i)).then(|| b.value(i));
            if av != bv {
                changed[i] = true;
            }
        }
    }
    for c in changed.iter_mut() {
        *c = !*c;
    }
    Ok(changed)
}

/// Apply `noise_type` to `orig`, retrying with fresh randomness (merging in
/// only the still-unchanged rows via `zip`) up to `MAX_NOISE_ATTEMPTS` times,
/// until every row's target columns differ from `orig` or attempts run out.
/// Returns the resulting batch, the final per-row "unchanged" mask (`true` =
/// still byte-identical to `orig` on this noise_type's target columns), and
/// the `target_cols` this noise_type actually touched.
///
/// When `target_cols` comes back empty (no column in this entity's schema
/// matches `noise_type` at all), every row is returned unchanged (`true`)
/// with no retries attempted — retrying can't manufacture a matching column
/// that doesn't exist; this is a schema/config-level gap, not a per-row
/// randomness one.
///
/// The retry merge only `zip()`s columns in `target_cols`:
/// `apply_noise_to_batch` already guarantees every other
/// column of `retried`/`noisy` is the same `Arc` as `orig`'s (untouched),
/// so `zip`-ing them is a provably-no-op pass over the full column just to
/// reconstruct a value that's already sitting there — wasted work that
/// scales with the entity's *total* column count instead of the handful
/// this noise_type actually targets. Confirmed empirically: for a 23-column
/// entity (kyc-sized) vs an 8-column one, this loop alone cost ~1.7x more
/// wall time despite touching the same single target column either way.
fn apply_noise_with_retry(
    orig: &RecordBatch,
    noise_type: &str,
    plan_cols: &[String],
    rng: &mut Rng,
    exclude_cols: &[String],
) -> Result<(RecordBatch, Vec<bool>, Vec<String>), String> {
    const MAX_NOISE_ATTEMPTS: usize = 4;
    let (mut noisy, target_cols) =
        apply_noise_to_batch(orig, noise_type, plan_cols, rng, exclude_cols)?;
    if target_cols.is_empty() {
        return Ok((noisy, vec![true; orig.num_rows()], target_cols));
    }
    let schema = noisy.schema();
    let target_idxs: Vec<usize> = target_cols
        .iter()
        .filter_map(|c| schema.index_of(c).ok())
        .collect();
    let mut unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
    let mut attempt = 1;
    while unchanged.iter().any(|&u| u) && attempt < MAX_NOISE_ATTEMPTS {
        let mut retry_rng = Rng::new(rng.next_u64());
        let (retried, _) =
            apply_noise_to_batch(orig, noise_type, plan_cols, &mut retry_rng, exclude_cols)?;
        let mask = arrow::array::BooleanArray::from_iter(unchanged.iter().map(|&u| Some(u)));
        let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
        for &i in &target_idxs {
            merged_cols[i] =
                arrow::compute::kernels::zip::zip(&mask, retried.column(i), noisy.column(i))
                    .map_err(|e| format!("zip noise retry col {i}: {e}"))?;
        }
        noisy = RecordBatch::try_new(noisy.schema(), merged_cols)
            .map_err(|e| format!("rebuild retried dup batch: {e}"))?;
        // Only re-test rows the zip above could actually have touched
        // — every row with `unchanged[i] == false`
        // is provably still `false` after this merge, see
        // `unchanged_row_mask_within`'s doc comment.
        unchanged = unchanged_row_mask_within(orig, &noisy, &target_cols, &unchanged)?;
        attempt += 1;
    }
    Ok((noisy, unchanged, target_cols))
}

/// Applies one additional noise pass, drawing an independent noise_type per
/// ROW from `all_types` instead of one type for the entire batch.
///
/// The original implementation drew a single `extra_type` for the whole
/// batch (`all_types[rng.next_usize(all_types.len())]`), applied once. If
/// that one draw landed on a type with no matching column in this entity's
/// schema at all (`noise_type_targets_column` false for every column —
/// common: `dates`/`identifiers`/`addresses`/`companies` are frequently
/// no-ops for entities that simply have no date/identifier/address/company
/// column, e.g. an aviation `passenger` row), `apply_noise_with_retry`
/// returns every row unchanged with *no retry attempted* (retrying can't
/// manufacture a matching column that doesn't exist) — so the entire batch
/// silently skipped this pass. With only 2-3 total passes, whether a batch's
/// "still identical after all passes" fate landed on such a no-op type was
/// close to a coin flip per pass, and that coin flip's odds shifted with
/// unrelated RNG-stream position (traced empirically: the same aviation
/// domain/seed/difficulty produced a real byte-identical-duplicate rate
/// swinging from ~12% to ~0.003% between two dataset sizes differing by
/// under 5%, tracking exactly when a per-entity batch boundary was crossed
/// elsewhere in the run — same root-cause family as `distribute_by_weight`).
///
/// Per-row assignment means every batch always has a healthy mix of types,
/// so a handful of schema-inapplicable categories can no longer determine
/// the whole batch's fate — same fix class as `noise::apply_random_subtype`,
/// applied one level up (noise *type* selection, not just sub-type selection
/// within one already-chosen type).
fn apply_extra_pass_per_row(
    noisy: &RecordBatch,
    unchanged: &mut [bool],
    all_types: &[&str],
    rng: &mut Rng,
    exclude_cols: &[String],
) -> Result<RecordBatch, String> {
    let n = noisy.num_rows();
    if all_types.is_empty() || n == 0 {
        return Ok(noisy.clone());
    }
    let per_row_type: Vec<usize> = (0..n).map(|_| rng.next_usize(all_types.len())).collect();

    // Partition rows by assigned type FIRST,
    // same technique as `noise::apply_random_subtype` — each row needs
    // noise from exactly one type, but the previous implementation called
    // `apply_noise_with_retry` on the FULL n-row batch once per active
    // type and threw away the n - |group| rows that mask didn't select:
    // O(len(all_types) * n) noise computation instead of O(n). On a schema
    // where every category has real target columns (e.g. kyc, all 9 hell
    // categories active), that's up to 9x redundant work per
    // extra pass. Grouping first means each row's noise is computed once,
    // on a batch sized to just its own group.
    // `row_group`/`row_pos` record, per original row, which type-group it
    // landed in and its position within that group's `idxs` (built in the
    // same pass as `groups`) — same technique as
    // `noise::apply_random_subtype` (transposed one level up): every noised value is read straight out of its group's
    // Arrow buffer at scatter time below, instead of round-tripping through
    // a `HashMap<String, Vec<Option<String>>>` of owned `String`s that a
    // final `StringArray::from` would then copy AGAIN.
    let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
    let mut row_group: Vec<u32> = Vec::with_capacity(n);
    let mut row_pos: Vec<u32> = Vec::with_capacity(n);
    for (i, &t) in per_row_type.iter().enumerate() {
        row_pos.push(groups[t].len() as u32);
        row_group.push(t as u32);
        groups[t].push(i as u32);
    }

    let schema = noisy.schema();
    let mut pass_unchanged = vec![true; n];
    let mut noised_by_group: Vec<Option<RecordBatch>> = vec![None; all_types.len()];
    // Per group, the column indices its noise_type actually targeted —
    // needed at scatter time to tell "this row's group touched this
    // column" (read `noised_by_group`) from "this row's group left this
    // column alone" (read `noisy`, unchanged) — different types touch
    // different column sets.
    let mut group_target_idxs: Vec<Vec<usize>> = vec![Vec::new(); all_types.len()];
    let mut touched_cols: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();

    for (type_idx, idxs) in groups.into_iter().enumerate() {
        if idxs.is_empty() {
            continue;
        }
        let ttype = all_types[type_idx];
        let idx_arr = UInt64Array::from_iter_values(idxs.iter().map(|&i| i as u64));
        let subset = pick_rows(noisy, &idx_arr)?;
        let (noised_subset, subset_unchanged, target_cols) =
            apply_noise_with_retry(&subset, ttype, &[], rng, exclude_cols)?;
        let target_idxs: Vec<usize> = target_cols
            .iter()
            .filter_map(|c| schema.index_of(c).ok())
            .collect();
        touched_cols.extend(target_idxs.iter().copied());
        group_target_idxs[type_idx] = target_idxs;
        for (local_pos, &orig_i) in idxs.iter().enumerate() {
            pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
        }
        noised_by_group[type_idx] = Some(noised_subset);
    }

    let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
    for col_idx in touched_cols {
        let orig_col = noisy.column(col_idx).as_string::<i32>();
        // Capacity sized to the column's actual existing value-buffer width
        // instead of a flat 16-byte guess (same motif as
        // `pool_lookup::pool_values`) — `orig_col.value_data()` is
        // the whole values buffer already sitting in memory for this
        // column, so its length is an exact, free-to-read bound.
        let cap_bytes = orig_col.value_data().len().max(n);
        let mut builder = StringBuilder::with_capacity(n, cap_bytes);
        // Per-column, not per-row: whether this
        // group even touches `col_idx`, and which noised `StringArray` to
        // read from it if so, depends only on `(g, col_idx)` — at most
        // `all_types.len()` distinct outcomes — but the previous loop
        // re-evaluated a linear `Vec::contains` scan AND an Arrow
        // `as_string::<i32>()` downcast (a dynamic type check) on every one
        // of the `n` rows. Precomputing per group here turns the inner loop
        // into a plain indexed lookup + branch.
        let per_group: Vec<Option<&arrow::array::StringArray>> = (0..group_target_idxs.len())
            .map(|g| {
                if group_target_idxs[g].contains(&col_idx) {
                    noised_by_group[g]
                        .as_ref()
                        .map(|b| b.column(col_idx).as_string::<i32>())
                } else {
                    None
                }
            })
            .collect();
        for i in 0..n {
            let g = row_group[i] as usize;
            match per_group[g] {
                Some(noised_col) => {
                    let pos = row_pos[i] as usize;
                    if noised_col.is_null(pos) {
                        builder.append_null();
                    } else {
                        builder.append_value(noised_col.value(pos));
                    }
                }
                None => {
                    if orig_col.is_null(i) {
                        builder.append_null();
                    } else {
                        builder.append_value(orig_col.value(i));
                    }
                }
            }
        }
        merged_cols[col_idx] = Arc::new(builder.finish());
    }
    let current = RecordBatch::try_new(noisy.schema(), merged_cols)
        .map_err(|e| format!("rebuild extra-pass batch: {e}"))?;

    for (u, pu) in unchanged.iter_mut().zip(pass_unchanged.iter()) {
        *u = *u && *pu;
    }
    Ok(current)
}

// ── Topological sort for FK dependency ordering ────────────────────────────

/// Order entity plans so that FK target entities are processed before dependents.
fn topological_sort(
    plans: &[EntityPlan],
    name_to_idx: &HashMap<&str, usize>,
) -> Result<Vec<usize>, String> {
    let n = plans.len();
    let mut in_degree = vec![0usize; n];
    let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
    for i in 0..n {
        for remap in &plans[i].fk_remaps {
            let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
                continue;
            };
            if target == i {
                continue;
            }
            adj[target].push(i); // target → i means target must come before i
            in_degree[i] += 1;
        }
    }
    let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
    let mut order = Vec::with_capacity(n);
    let mut idx = 0;
    while idx < queue.len() {
        let v = queue[idx];
        idx += 1;
        order.push(v);
        for &next in &adj[v] {
            in_degree[next] -= 1;
            if in_degree[next] == 0 {
                queue.push(next);
            }
        }
    }
    if order.len() != n {
        return Err("FK dependency cycle detected among entities".into());
    }
    Ok(order)
}

/// Dataset sink: writes directly in the requested output format, avoiding
/// an uncompressed IPC intermediate + full reread when `parquet` is
/// requested (mirrors the `GtSink` pattern in `gt.rs`).
pub(crate) enum DatasetWriter {
    Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
    Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
}

impl DatasetWriter {
    fn new(
        output_format: &str,
        path: &str,
        schema: &Arc<Schema>,
        metadata: &RunMetadata,
    ) -> Result<Self, String> {
        let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
        if output_format == "parquet" {
            use parquet::basic::{Compression, ZstdLevel};
            use parquet::file::properties::WriterProperties;
            let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
            let meta_kv = parquet_kv(metadata);
            let props = WriterProperties::builder()
                .set_compression(Compression::ZSTD(zstd))
                // 1M rows/group (parquet-rs' own default) bounds the
                // resident, uncompressed column-chunk buffers to one row
                // group's worth at a time instead of the whole dataset.
                .set_max_row_group_row_count(Some(1_000_000))
                .set_key_value_metadata(Some(meta_kv))
                .build();
            let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
                .map_err(|e| format!("ArrowWriter error: {e}"))?;
            Ok(DatasetWriter::Parquet(Box::new(writer)))
        } else {
            let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
                .map_err(|e| format!("IPC FileWriter error: {e}"))?;
            Ok(DatasetWriter::Ipc(Box::new(writer)))
        }
    }

    pub(crate) fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
        match self {
            DatasetWriter::Ipc(w) => w.write(batch).map_err(|e| format!("write ipc: {e}")),
            DatasetWriter::Parquet(w) => w.write(batch).map_err(|e| format!("write parquet: {e}")),
        }
    }

    fn finish(self) -> Result<(), String> {
        match self {
            DatasetWriter::Ipc(mut w) => w.finish().map_err(|e| format!("finish ipc: {e}")),
            DatasetWriter::Parquet(w) => w
                .close()
                .map(|_| ())
                .map_err(|e| format!("close parquet: {e}")),
        }
    }
}

// ── RSS-by-phase debug instrumentation ──────────────────────────────────────
//
// Non-invasive: only samples/prints when `RUST_LOG=debug` (or
// lower) is enabled -- default runs pay one `log::log_enabled!` check per
// call site and nothing else. Reuses `sysinfo` (already a normal
// dependency, already used for the same purpose in `main.rs`'s
// `warn_if_memory_tight`) instead of adding a new one for a debug
// probe. Not a substitute for real allocation-site attribution --
// just phase-boundary RSS deltas to localize which part of the pipeline
// dominates peak RAM before writing an isolated micro-bench for it.
fn log_rss(label: &str) {
    if !log::log_enabled!(log::Level::Debug) {
        return;
    }
    let pid = sysinfo::Pid::from_u32(std::process::id());
    let Ok(mut sys) = sysinfo::System::new() else {
        return;
    };
    sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
    if let Some(p) = sys.process(pid) {
        log::debug!("[rss] {label}: rss={:.1} MB", p.memory() as f64 / 1e6);
    }
}

// ── Main pipeline ──────────────────────────────────────────────────────────

pub fn run_pipeline(
    ctx: &Context,
    config: &PipelineConfig,
    output_dir: &str,
) -> Result<PipelineOutput, String> {
    run_pipeline_with_progress(ctx, config, output_dir, None)
}

/// Same as [`run_pipeline`], but invokes `progress(records_written_so_far,
/// approx_total)` after every base-record batch is written to the dataset
/// file — `approx_total` is `config.size` (the requested record count; the
/// actual total including duplicates/HN/canary rows differs slightly, close
/// enough for a progress display). Callers that don't need progress
/// reporting should use [`run_pipeline`] instead.
pub fn run_pipeline_with_progress(
    ctx: &Context,
    config: &PipelineConfig,
    output_dir: &str,
    progress: Option<&mut dyn FnMut(usize, usize)>,
) -> Result<PipelineOutput, String> {
    run_pipeline_chunked(ctx, config, output_dir, progress, None).map(|(out, _)| out)
}

/// Same as [`run_pipeline_with_progress`], but threads `record_id`/
/// `master_id` counters through `offsets` (`None` == all-zero, identical to
/// a normal single run) and returns the counters' final values alongside
/// the usual output — the mechanism `--chunk-size` orchestration uses to
/// keep IDs globally contiguous across a sequence of otherwise-independent
/// chunk runs. See `ChunkOffsets` for why this exists.
pub fn run_pipeline_chunked(
    ctx: &Context,
    config: &PipelineConfig,
    output_dir: &str,
    mut progress: Option<&mut dyn FnMut(usize, usize)>,
    offsets: Option<&ChunkOffsets>,
) -> Result<(PipelineOutput, ChunkOffsets), String> {
    let offsets = offsets.cloned().unwrap_or_default();
    std::fs::create_dir_all(output_dir)
        .map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
    let t_start = std::time::Instant::now();

    let t_alloc = t_start.elapsed().as_secs_f64();
    log_rss("start");

    // ── Phase 14: Streaming pipeline ─────────────────────────────────────
    // Generates, extracts FK/HN, remaps, and sinks each batch per entity
    // without storing all entity batches in memory.
    let t1e = std::time::Instant::now();

    // Metadata (including the `dupehell.timestamp` snapshot) is computed
    // once here and reused for every output file (dataset, graph, GT,
    // parquet conversions) — computing it per-file let `SystemTime::now()`
    // drift across a Unix-second boundary on longer runs, making
    // `dupehell.timestamp` inconsistent between files from the same run.
    let metadata = build_metadata_map(config);
    // Build the full schema once (union of all entity columns + metadata)
    let full_arc = Arc::new(build_full_schema(config, &metadata));
    // One shared dictionary per low-cardinality column, built from the full
    // known value set before any batch is written — every
    // `domain`/`entity_type` array for the rest of this file reuses these
    // same `values`, required for IPC dictionary consistency (see
    // `DictValues`). `entity_type`'s value set is exactly the written
    // (non-pool-only) entity plans' names — base/dup/HN/canary rows never
    // carry any other entity_type value.
    let domain_dict = DictValues::new([config.domain.clone()]);
    let entity_type_dict = DictValues::new(
        config
            .entity_plans
            .iter()
            .filter(|p| !p.pool_only)
            .map(|p| p.name.clone()),
    );
    let dataset_ext = if config.output_format == "parquet" {
        "parquet"
    } else {
        "ipc"
    };
    let dataset_path = format!("{}/{}.{}", output_dir, config.run_id, dataset_ext);
    let mut writer =
        DatasetWriter::new(&config.output_format, &dataset_path, &full_arc, &metadata)?;

    // Topological sort for FK dependency processing order
    let name_to_idx: HashMap<&str, usize> = config
        .entity_plans
        .iter()
        .enumerate()
        .map(|(i, p)| (p.name.as_str(), i))
        .collect();
    let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;

    // Cross-entity storage (lightweight)
    let mut fk_pools: FkPoolMap = HashMap::new();
    let mut hn_pools: HashMap<String, HnPool> = HashMap::new();

    // GT streaming: each entity batch is classified and written to the
    // ground-truth file as soon as its duplicates are known (see
    // `crate::gt::GtAccumulator`).
    let gt_ext = if config.output_format == "parquet" {
        "parquet"
    } else {
        "ipc"
    };
    let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
    let mut gt_acc: Option<crate::gt::GtAccumulator> = if config.skip_ground_truth {
        None
    } else {
        Some(crate::gt::GtAccumulator::new(
            config.difficulty.as_str(),
            &config.output_format,
            &gt_path,
            &metadata,
            config.graph_enabled,
        )?)
    };

    let mut global_rid_offset = offsets.rid_offset;

    // Stops the run cleanly if the output volume runs low on space (polled
    // once per batch, see `disk::DiskGuard`); `None` when the volume can't
    // be identified.
    let mut disk_guard = crate::disk::DiskGuard::new(std::path::Path::new(output_dir));

    // ── Graph output (opt-in via --graph) ───────────────────────────────
    const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
    let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
    let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
    let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
    let mut graph_nodes_path: Option<String> = None;
    let mut graph_edges_path: Option<String> = None;
    if config.graph_enabled {
        let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
        let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
        node_writer = Some(
            crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
                .map_err(|e| format!("init node writer: {e}"))?,
        );
        edge_writer = Some(
            crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata, subtype_dict(config))
                .map_err(|e| format!("init edge writer: {e}"))?,
        );
        graph_nodes_path = Some(nodes_ipc);
        graph_edges_path = Some(edges_ipc);
    }

    let mut _t_write = 0.0f64;
    let mut _t_meta = 0.0f64;
    let mut _t_dup = 0.0f64;
    let mut _write_calls = 0u64;
    // Compute HN pool size: need at least 2× the max hn_cfg.count for any entity type
    let hn_pool_sizing: std::collections::HashMap<&str, usize> =
        config
            .hard_neg_types
            .iter()
            .fold(std::collections::HashMap::new(), |mut m, hn| {
                let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
                *entry = (*entry).max(hn.count);
                m
            });
    let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
    // FK remap (fk_remap.rs) samples uniformly at random from this pool via
    // `take` — it has no need to see every identifier the entity ever
    // produced, only a large-enough sample for the target diversity. Capping
    // it (same idea as `hn_max_pool` above) keeps its RAM bounded instead of
    // O(n_base) per entity with an identifier column. (`FK_POOL_CAP` is now
    // module-level, see its doc comment.)

    // Entities whose FK pool is ever read (`fk_pools.get(&remap.target_entity)`
    // below, and the identical lookup in `canary::generate_all`, which reuses
    // the same `fk_remaps`). Extracting identifiers for an entity outside
    // this set is pure waste — the pool built for it is never consulted.
    // Entities present only as an FK identifier pool (`--only-entity`),
    // never written as graph nodes. A `fk`-typed edge pointing at one of
    // these would dangle (target node absent from the nodes output), so
    // graph edges are suppressed for remaps into this set below.
    let pool_only_names: std::collections::HashSet<&str> = config
        .entity_plans
        .iter()
        .filter(|p| p.pool_only)
        .map(|p| p.name.as_str())
        .collect();

    let fk_targets: std::collections::HashSet<&str> = config
        .entity_plans
        .iter()
        .flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
        .collect();

    for &plan_idx in &entity_order {
        let plan = &config.entity_plans[plan_idx];
        if plan.n_base == 0 {
            continue;
        }

        let prefix = entity_prefix(plan_idx);
        let master_base = offsets.master_base(plan_idx);
        let id_base = offsets.id_base(plan_idx);
        // Parsed once per entity, not once per batch — see
        // `entity_gen::parse_columns`'s doc comment.
        let parsed_columns = crate::entity_gen::parse_columns(&plan.columns_json)?;

        // Per-entity streaming state
        let fk_targeted = fk_targets.contains(plan.name.as_str());
        // FK pool: a uniform reservoir sample (Algorithm R) of at most
        // FK_POOL_CAP of this entity's identifiers, with their record index
        // for graph edges. It used to keep the *first* FK_POOL_CAP
        // identifiers, so past that many base rows no later record of the
        // entity was ever referenced by another entity's FK.
        let mut fk_ids: Vec<String> = Vec::new();
        let mut fk_rids: Vec<u64> = Vec::new();
        let mut fk_seen: usize = 0;
        let mut fk_pool_rng = Rng::new(crate::rng::derive_seed(
            config.seed,
            &[
                crate::rng::stream::FK_POOL,
                crate::rng::name_key(&plan.name),
            ],
        ));
        // Only entities actually referenced by a `hard_neg_types` entry ever
        // have their HN pool read (see `hn_pools.get(&hn_cfg.entity_type)`
        // below); accumulating slices for the others just pins whole-batch
        // Arrow buffers in `hn_pools` for the rest of the run for nothing.
        let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
        let mut hn_slices: Vec<RecordBatch> = Vec::new();
        let mut hn_slice_rids: Vec<ArrayRef> = Vec::new();
        // Per-entity streams (see `rng::derive_seed`): with a shared seed,
        // the i-th row of every entity drew the same duplicates/FK targets.
        let entity_key = crate::rng::name_key(&plan.name);
        let mut batch_rng = Rng::new(crate::rng::derive_seed(
            config.seed,
            &[crate::rng::stream::DUPLICATES, entity_key],
        ));
        let mut fk_rng = Rng::new(crate::rng::derive_seed(
            config.seed,
            &[crate::rng::stream::FK_REMAP, entity_key],
        ));

        // Each noise_type's planned `count` is spread across every batch of
        // masters in proportion to batch size (largest-remainder method),
        // so duplicates land on masters throughout the whole entity instead
        // of being sampled only from whichever batch happens to run last —
        // see the "Dups" section below, at the end of each batch iteration.
        let batch_bounds: Vec<(usize, usize)> = (0..plan.n_base)
            .step_by(crate::entity_gen::BATCH_SIZE)
            .map(|offset| {
                (
                    offset,
                    (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE),
                )
            })
            .collect();
        let batch_weights: Vec<usize> = batch_bounds.iter().map(|&(_, bn)| bn).collect();
        let per_batch_noise_counts: Vec<Vec<usize>> = plan
            .noise_types
            .iter()
            .map(|n| distribute_by_weight(n.count, &batch_weights))
            .collect();
        // Both invariant per entity, not per batch
        // — hoisted out of the "Dups" section below, alongside this
        // entity's other precomputed invariants (`batch_bounds`,
        // `per_batch_noise_counts`), instead of being reconstructed inside
        // `if has_dups { … }` on every batch that has duplicates.
        let fk_exclude_cols: Vec<String> = plan
            .fk_remaps
            .iter()
            .map(|r| r.source_col.clone())
            .collect();
        let all_types: Vec<&str> = plan
            .noise_types
            .iter()
            .map(|n| n.noise_type.as_str())
            .collect();

        // Build col_lookup from first batch (it has the entity's schema)
        let mut col_lookup: Option<Vec<Option<usize>>> = None;
        // Phase 12.3: null cache per entity plan
        let mut null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
        let mut const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();

        // Stream: generate → FK remap → FK extract → HN sample → IPC write → GT collect
        for (batch_idx, offset) in (0..plan.n_base)
            .step_by(crate::entity_gen::BATCH_SIZE)
            .enumerate()
        {
            let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
            let batch_seed = crate::rng::derive_seed(
                config.seed,
                &[
                    crate::rng::stream::ENTITY_BATCH,
                    entity_key,
                    batch_idx as u64,
                ],
            );

            // Generate batch
            let rb = crate::entity_gen::generate_entity_batch_parsed(
                ctx,
                &parsed_columns,
                batch_n,
                batch_seed,
                id_base + offset,
            )?;

            // FK remap (uses pools from previously processed entities)
            let mut fk_edges_by_remap: Vec<(String, ArrayRef, f64)> = Vec::new();
            let rb = if !plan.fk_remaps.is_empty() {
                let mut r = rb;
                for remap in &plan.fk_remaps {
                    if let Some(pool) = fk_pools.get(&remap.target_entity) {
                        // Suppress edge emission into a pool-only target —
                        // its rows were never written as graph nodes, so an
                        // edge to them would dangle (see `pool_only_names`).
                        let want_edge = config.graph_enabled
                            && !pool_only_names.contains(remap.target_entity.as_str());
                        let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
                            &r,
                            pool,
                            &remap.source_col,
                            &mut fk_rng,
                            want_edge,
                        )?;
                        r = remapped;
                        if want_edge && let Some(rids) = target_rids {
                            fk_edges_by_remap.push((
                                remap.source_col.clone(),
                                rids,
                                remap.fk_weight,
                            ));
                        }
                    }
                }
                r
            } else {
                rb
            };

            // Feed this batch's identifiers to the FK pool reservoir (see
            // `fk_ids` above).
            if fk_targeted
                && let Some(ref id_col) = plan.identifier_col
                && let Some(col) = rb.column_by_name(id_col)
            {
                let s = col.as_string::<i32>();
                for i in 0..s.len() {
                    if s.is_null(i) {
                        continue;
                    }
                    let rid = (global_rid_offset + i) as u64;
                    if fk_seen < FK_POOL_CAP {
                        fk_ids.push(s.value(i).to_string());
                        fk_rids.push(rid);
                    } else {
                        let j = fk_pool_rng.next_usize(fk_seen + 1);
                        if j < FK_POOL_CAP {
                            fk_ids[j].clear();
                            fk_ids[j].push_str(s.value(i));
                            fk_rids[j] = rid;
                        }
                    }
                    fk_seen += 1;
                }
            }

            // Pool-only entity (`--only-entity`): its rows exist solely to
            // seed `fk_pools` (extracted just above) — no noise, no writing,
            // no GT, no graph nodes, and `global_rid_offset` does not
            // advance for it (record_ids stay contiguous across the rows
            // that ARE written).
            if plan.pool_only {
                continue;
            }

            // HN pool: accumulate slices up to max pool size
            let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
            if hn_targeted && hn_accum < hn_max_pool {
                let take = (hn_max_pool - hn_accum).min(batch_n);
                hn_slices.push(rb.slice(0, take));
                // Capture the global record_ids for these sliced rows so the
                // HN pool can later resolve `hard_neg` edge targets. Use
                // `global_rid_offset`, never the per-entity `offset`.
                if config.graph_enabled {
                    hn_slice_rids
                        .push(record_id_array(global_rid_offset..global_rid_offset + take));
                }
            }

            // Build col_lookup from first batch's schema
            if col_lookup.is_none() {
                let plan_schema = rb.schema();
                col_lookup = Some(
                    full_arc
                        .fields()
                        .iter()
                        .skip(4)
                        .map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
                        .collect(),
                );
            }

            // Master IDs for this batch
            let batch_mid_arr = master_id_array_range(&prefix, master_base + offset, batch_n);

            // Add metadata + IPC write
            let rid_arr = record_id_array(global_rid_offset..global_rid_offset + batch_n);
            let t_m0 = std::time::Instant::now();
            let base_rb = add_metadata_and_align(
                &rb,
                &config.domain,
                &plan.name,
                &domain_dict,
                &entity_type_dict,
                &rid_arr,
                &batch_mid_arr,
                &full_arc,
                col_lookup.as_ref().unwrap(),
                &mut null_cache,
                &mut const_arr_cache,
            )?;
            _t_meta += t_m0.elapsed().as_secs_f64();

            // Graph: write base nodes + FK edges (only when --graph)
            if config.graph_enabled
                && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
            {
                nw.write_batch(&base_rb)
                    .map_err(|e| format!("write node: {e}"))?;
                // `rid_arr` and each remap's `target_rids`
                // are already whole `StringArray`s in row order -- push them
                // as a single batch (one `Arc` clone each) instead of the
                // previous `for i in 0..batch_n` scalar loop that
                // re-extracted `.value(i)` and re-appended through a
                // `StringBuilder` plus a `HashMap` lookup per row.
                for (subtype, target_rids, weight) in &fk_edges_by_remap {
                    ew.push_batch(&rid_arr, target_rids, "fk", subtype, *weight)
                        .map_err(|e| format!("push fk edge batch: {e}"))?;
                }
            }

            let t_w0 = std::time::Instant::now();
            writer
                .write(&base_rb)
                .map_err(|e| format!("write base: {e}"))?;
            _t_write += t_w0.elapsed().as_secs_f64();
            _write_calls += 1;

            // GT for this batch is written once its duplicates (below) are
            // known — they're drawn from this batch's masters only.
            let base_first_rid = global_rid_offset;
            // Duplicate rows of this batch for GT: (rows, global master
            // index per row, is_identical per row, first record index).
            let mut dup_gt: Option<(RecordBatch, Vec<usize>, Vec<bool>, usize)> = None;

            global_rid_offset += batch_n;
            if let Some(cb) = &mut progress {
                cb(global_rid_offset, config.size);
            }

            // ── Dups ─────────────────────────────────────────────────────
            // Sampled from *this* batch's masters only (`rb`/`batch_n`/
            // `offset`), using this batch's share of each noise_type's
            // total count (`per_batch_noise_counts`, computed once before
            // the loop). Every batch of masters gets duplicates this way,
            // instead of concentrating all of them on whichever batch
            // happened to run last — see `distribute_by_weight`.
            let has_dups: bool = per_batch_noise_counts
                .iter()
                .any(|counts| counts[batch_idx] > 0);
            if has_dups {
                let t_d0 = std::time::Instant::now();
                let mut dup_batches: Vec<RecordBatch> = Vec::new();
                let mut dup_master_idx_buf: Vec<usize> = Vec::new();

                // Pre-generate indices + parallel noise (Phase 13c)
                let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
                    .noise_types
                    .iter()
                    .enumerate()
                    .filter_map(|(ni, n)| {
                        let cnt = per_batch_noise_counts[ni][batch_idx];
                        if cnt == 0 {
                            return None;
                        }
                        let mut b = UInt64Builder::with_capacity(cnt);
                        for _ in 0..cnt {
                            b.append_value(batch_rng.next_usize(batch_n) as u64);
                        }
                        Some((
                            b.finish(),
                            batch_rng.next_u64(),
                            n.noise_type.as_str(),
                            n.columns.as_slice(),
                            cnt,
                        ))
                    })
                    .collect();

                // Distinct active noise_type names for this entity (hoisted
                // above, see the H8 comment), used by additional noise
                // passes (below) to independently draw a type to apply on
                // top of the first pass's result.
                let all_types_ref = &all_types;
                let noise_passes = config.noise_passes.max(1);

                type DupNoiseResult = Result<(RecordBatch, Vec<usize>, Vec<bool>), String>;
                let fk_exclude_ref = &fk_exclude_cols;
                let rb_ref = &rb;
                // `rayon::scope`'s `par_iter().map(...).collect()` runs each
                // noise-type bucket on Rayon's already-warm global pool
                // instead of spawning a fresh OS thread per bucket per batch
                // (the previous `std::thread::scope` pattern) — an isolated
                // micro-benchmark (`examples/profile_thread_spawn.rs`)
                // measured raw OS thread spawn/join at 12-19x the cost of the
                // pooled equivalent at this spawn count/workload shape.
                // `par_iter` preserves input order in its output (unlike a
                // channel), so `results` ends up in the same `ndata` order as
                // before — required downstream since `dup_mids_buf`/
                // `global_rid_offset` assignment depends on that order, not on
                // which bucket happens to finish first.
                let results: Vec<DupNoiseResult> = ndata
                    .par_iter()
                    .map(|(indices, seed, ntype, cols, cnt)| {
                        let dup = pick_rows(rb_ref, indices)?;
                        let mut rng = Rng::new(*seed);

                        // First pass: the noise_type this bucket was
                        // assigned, with the retry-until-changed guarantee
                        // (see `apply_noise_with_retry`) so a "noised"
                        // duplicate can't silently stay a byte-for-byte
                        // copy of its master.
                        let (mut noisy, mut unchanged, _) =
                            apply_noise_with_retry(&dup, ntype, cols, &mut rng, fk_exclude_ref)?;

                        // Additional independent passes (difficulty-
                        // controlled, see `schema::DifficultySettings::
                        // passes`): each draws its own noise_type from this
                        // entity's active list and applies it on top of the
                        // previous pass's result, compounding corruption.
                        // This — not adding more categories to
                        // `noise_types` — is what makes a tier reliably
                        // harder without diluting any single category's
                        // weight (see the doc comment on `passes`). A row
                        // is only "still identical" overall once every
                        // pass left it unchanged.
                        for _ in 1..noise_passes {
                            noisy = apply_extra_pass_per_row(
                                &noisy,
                                &mut unchanged,
                                all_types_ref,
                                &mut rng,
                                fk_exclude_ref,
                            )?;
                        }

                        let mut mb = Vec::with_capacity(*cnt);
                        for j in 0..*cnt {
                            // `indices` are local to `rb` (0..batch_n); the master_id
                            // is a pure function of (prefix, global index) — same
                            // formula as `batch_mids` above — so it's recomputed here
                            // instead of being cloned out of a retained master_id_pool.
                            let global_idx = master_base + offset + indices.value(j) as usize;
                            mb.push(global_idx);
                        }
                        Ok((noisy, mb, unchanged))
                    })
                    .collect();
                let mut dup_is_identical_buf: Vec<bool> = Vec::new();
                for res in results {
                    let (rb, mb, unchanged) = res?;
                    dup_batches.push(rb);
                    dup_master_idx_buf.extend(mb);
                    dup_is_identical_buf.extend(unchanged);
                }
                _t_dup += t_d0.elapsed().as_secs_f64();

                // Write dups
                if !dup_batches.is_empty() {
                    let dup_total = dup_master_idx_buf.len();
                    let dup_rid_arr =
                        record_id_array(global_rid_offset..global_rid_offset + dup_total);
                    let dup_mid_arr = master_id_array_from_indices(
                        &prefix,
                        dup_master_idx_buf.iter().copied(),
                        dup_total,
                    );

                    let concated = if dup_batches.len() == 1 {
                        dup_batches.into_iter().next().unwrap()
                    } else {
                        let schema = dup_batches[0].schema();
                        let n_fields = dup_batches[0].num_columns();
                        let mut concat_arrays = Vec::with_capacity(n_fields);
                        for i in 0..n_fields {
                            let refs: Vec<&dyn Array> =
                                dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
                            concat_arrays.push(
                                arrow::compute::concat(&refs)
                                    .map_err(|e| format!("concat dup col {i}: {e}"))?,
                            );
                        }
                        RecordBatch::try_new(schema, concat_arrays)
                            .map_err(|e| format!("concat dups: {e}"))?
                    };

                    let dup_rb_full = add_metadata_and_align(
                        &concated,
                        &config.domain,
                        &plan.name,
                        &domain_dict,
                        &entity_type_dict,
                        &dup_rid_arr,
                        &dup_mid_arr,
                        &full_arc,
                        col_lookup.as_ref().unwrap(),
                        &mut null_cache,
                        &mut const_arr_cache,
                    )?;

                    let t_wd = std::time::Instant::now();
                    writer
                        .write(&dup_rb_full)
                        .map_err(|e| format!("write dups: {e}"))?;
                    _t_write += t_wd.elapsed().as_secs_f64();
                    _write_calls += 1;

                    // Graph: write duplicate nodes (edges emitted post-GT)
                    if config.graph_enabled
                        && let Some(nw) = node_writer.as_mut()
                    {
                        nw.write_batch(&dup_rb_full)
                            .map_err(|e| format!("write dup node: {e}"))?;
                    }

                    if gt_acc.is_some() {
                        dup_gt = Some((
                            dup_rb_full,
                            dup_master_idx_buf,
                            dup_is_identical_buf,
                            global_rid_offset,
                        ));
                    }
                    global_rid_offset += dup_total;
                }
            }

            if let Some(acc) = gt_acc.as_mut() {
                let base = crate::gt::BaseRows {
                    record_ids: base_rb.column(0),
                    entity_types: base_rb.column(2),
                    master_ids: base_rb.column(3),
                    entity_idx: plan_idx as u64,
                    first_master_idx: (master_base + offset) as u64,
                    first_record_idx: base_first_rid as u64,
                };
                let dups =
                    dup_gt
                        .as_ref()
                        .map(|(rb, master_idx, ident, first_rid)| crate::gt::DupRows {
                            record_ids: rb.column(0),
                            entity_types: rb.column(2),
                            master_ids: rb.column(3),
                            master_idx,
                            is_identical: ident,
                            first_record_idx: *first_rid as u64,
                        });
                acc.push_entity_batch(&base, dups.as_ref())?;
            }
            if let Some(g) = disk_guard.as_mut() {
                g.check()?;
            }
        }

        // Save FK pool for cross-entity remapping
        if let Some(id_col) = &plan.identifier_col
            && !fk_ids.is_empty()
        {
            let arr = Arc::new(arrow::array::StringArray::from_iter_values(&fk_ids)) as ArrayRef;
            if config.graph_enabled {
                // Graph mode: also persist the target record_ids (column 1)
                // so FK edges can be emitted by later-source entities.
                let mut rid_builder = arrow::array::StringBuilder::with_capacity(
                    fk_rids.len(),
                    fk_rids.len() * RID_LEN,
                );
                let mut rid_buf: Vec<u8> = Vec::with_capacity(RID_LEN);
                for &rid in &fk_rids {
                    rid_buf.clear();
                    append_record_id(&mut rid_buf, rid);
                    rid_builder.append_value(
                        std::str::from_utf8(&rid_buf)
                            .expect("append_record_id writes only ASCII digits"),
                    );
                }
                let rid_arr = Arc::new(rid_builder.finish()) as ArrayRef;
                let schema = Arc::new(Schema::new(vec![
                    Field::new(id_col, DataType::Utf8, true),
                    Field::new("record_id", DataType::Utf8, true),
                ]));
                if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
                    fk_pools.insert(plan.name.clone(), pool_rb);
                }
            } else {
                let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
                if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
                    fk_pools.insert(plan.name.clone(), pool_rb);
                }
            }
        }

        // Build HN pool from collected slices
        if !hn_slices.is_empty() {
            let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
            let batch = if hn_slices.len() == 1 {
                // `rb.slice(0, take)` (pushed above) is zero-copy: it keeps
                // the *entire* source batch's buffers alive (up to
                // BATCH_SIZE=500_000 rows × all columns) for the rest of the
                // run just to serve `take` rows. Re-materializing through the
                // `take` kernel (same helper as `pick_rows` below) builds
                // compact, right-sized buffers instead.
                let only = hn_slices.into_iter().next().unwrap();
                let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
                pick_rows(&only, &idx)?
            } else {
                let schema = hn_slices[0].schema();
                let n_fields = schema.fields().len();
                let mut concat_arrays = Vec::with_capacity(n_fields);
                for i in 0..n_fields {
                    let refs: Vec<&dyn Array> =
                        hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
                    concat_arrays.push(
                        arrow::compute::concat(&refs)
                            .map_err(|e| format!("hn pool concat col {i}: {e}"))?,
                    );
                }
                RecordBatch::try_new(schema, concat_arrays)
                    .map_err(|e| format!("hn pool concat batch: {e}"))?
            };
            // `record_ids` aligns positionally with the concatenated pool rows.
            let record_ids: ArrayRef = if config.graph_enabled {
                let refs: Vec<&dyn Array> = hn_slice_rids.iter().map(|a| a.as_ref()).collect();
                arrow::compute::concat(&refs).map_err(|e| format!("hn rid concat: {e}"))?
            } else {
                Arc::new(StringArray::new_null(0))
            };
            hn_pools.insert(
                plan.name.clone(),
                HnPool {
                    batch,
                    total_count,
                    record_ids,
                },
            );
        }

        // Per-entity RSS checkpoint: the aggregate "after entity batches"
        // checkpoint below only samples once for the whole domain, which
        // can't tell apart an entity that dominates RAM from one that's
        // negligible when a domain has several entities.
        log_rss(&format!(
            "after entity '{}' ({} base rows)",
            plan.name, plan.n_base
        ));
    }
    let t1e_elapsed = t1e.elapsed().as_secs_f64();
    log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
    log_rss("after entity batches (fk_pools+hn_pools built)");

    // ── Phase 2: Hard negatives ────────────────────────────────────────────
    let t2 = std::time::Instant::now();
    // Phase 12.3: separate null cache for HN section (different entity types)
    let mut hn_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
    let mut hn_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
    // Global counter for HN master_ids (shared across every hard_neg_types
    // entry in this run) instead of a random draw from a fixed 10M space:
    // at tens/hundreds of thousands of HN rows, `next_usize(10_000_000)`
    // has a real birthday-paradox chance of two unrelated HN rows landing
    // on the same master_id, which `gt.rs` would then count as sharing a
    // master — silently mislabeling a hard negative as a match. A counter
    // is unique by construction, eliminating the collision entirely.
    let mut hn_master_id_counter: u64 = offsets.hn_master_counter;
    for (hn_idx, hn_cfg) in config.hard_neg_types.iter().enumerate() {
        if hn_cfg.count == 0 {
            continue;
        }
        let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
            continue;
        };
        let hn_total = pool_data.total_count;

        let pool_rb = &pool_data.batch;
        let n_hn_max = hn_cfg.count.min(hn_total / 2);
        if n_hn_max == 0 {
            continue;
        }

        let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
            pool_rb,
            &hn_cfg.config_json,
            n_hn_max,
            // One stream per hard_neg_types entry: a shared seed drew the
            // same pair indices for every type.
            crate::rng::derive_seed(config.seed, &[crate::rng::stream::HARD_NEG, hn_idx as u64]),
            config.graph_enabled,
        )?;

        let n_hn = hn_rb.num_rows();
        if n_hn == 0 {
            continue;
        }

        let hn_rid_arr = record_id_array(global_rid_offset..global_rid_offset + n_hn);
        let hn_mid_arr = hn_master_id_array(hn_master_id_counter, n_hn);
        hn_master_id_counter += n_hn as u64;

        // Build col lookup for the HN entity type schema
        let hn_schema = hn_rb.schema();
        let hn_col_lookup: Vec<Option<usize>> = full_arc
            .fields()
            .iter()
            .skip(4)
            .map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
            .collect();

        let hn_rb_full = add_metadata_and_align(
            &hn_rb,
            &config.domain,
            &hn_cfg.entity_type,
            &domain_dict,
            &entity_type_dict,
            &hn_rid_arr,
            &hn_mid_arr,
            &full_arc,
            &hn_col_lookup,
            &mut hn_null_cache,
            &mut hn_const_arr_cache,
        )?;

        let t_wh = std::time::Instant::now();
        writer
            .write(&hn_rb_full)
            .map_err(|e| format!("write hn: {e}"))?;
        _t_write += t_wh.elapsed().as_secs_f64();
        _write_calls += 1;

        // Graph: write HN nodes + hard_neg edges (source = HN rid,
        // target = the pool row referenced by idx_a)
        if config.graph_enabled
            && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
        {
            nw.write_batch(&hn_rb_full)
                .map_err(|e| format!("write hn node: {e}"))?;
            if let Some((idx_a, pattern)) = hn_src {
                // Gather the target ids with the `take`
                // kernel (vectorized) instead of a scalar
                // `for i in .. { pool_rid_arr.value(idx_a[i]) }` loop, then
                // push the whole edge batch in one call -- same reasoning as
                // the FK-edge site above.
                let take_n = n_hn.min(idx_a.len());
                let idx_arr = arrow::array::UInt64Array::from_iter_values(
                    idx_a.iter().take(take_n).map(|&i| i as u64),
                );
                let tgt_arr = arrow::compute::take(&pool_data.record_ids, &idx_arr, None)
                    .map_err(|e| format!("take hn target ids: {e}"))?;
                let src_arr = hn_rid_arr.slice(0, take_n);
                ew.push_batch(&src_arr, &tgt_arr, "hard_neg", &pattern, 1.0)
                    .map_err(|e| format!("push hn edge batch: {e}"))?;
            }
        }

        if let Some(acc) = gt_acc.as_mut() {
            acc.push_hard_neg_batch(
                hn_rb_full.column(0),
                hn_rb_full.column(2),
                hn_rb_full.column(3),
            )?;
        }

        global_rid_offset += n_hn;
    }
    let t2_elapsed = t2.elapsed().as_secs_f64();
    // `hn_pools` is never read again past this point, but its natural scope
    // end is the whole function — dropping it explicitly here frees its RAM
    // (measured 125.1 MB on a 50M-record aviation/hell run)
    // as soon as it's actually done being useful instead of letting it sit
    // resident through canary generation, the dataset writer's finish, and
    // GT accumulation.
    drop(hn_pools);
    log_rss("after hard negatives");

    // ── Canary records ────────────────────────────────────────────────────
    {
        let mut canary_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
        let mut canary_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
        crate::canary::generate_all(
            ctx,
            config,
            &full_arc,
            &domain_dict,
            &entity_type_dict,
            &mut canary_null_cache,
            &mut canary_const_arr_cache,
            &mut global_rid_offset,
            &fk_pools,
            &offsets,
            &mut writer,
            &mut node_writer,
            gt_acc.as_mut(),
        )?;
    }
    log_rss("after canary records");

    // ── Phase 3: Finalize + GT (streamed, see crate::gt::GtAccumulator) ────
    let t3 = std::time::Instant::now();
    writer
        .finish()
        .map_err(|e| format!("finish dataset writer: {e}"))?;
    let t3a_elapsed = t3.elapsed().as_secs_f64();
    log_rss("after dataset writer finish");

    let t3b = std::time::Instant::now();
    let t_gt0 = std::time::Instant::now();
    let crate::gt::GtResult {
        n_exact_dup,
        n_fuzzy_dup,
        n_hard_neg,
        n_unique,
        n_masters,
        cluster_map,
    } = match gt_acc {
        Some(acc) => acc.finish()?,
        // `--skip-ground-truth`: nothing was classified and no
        // `_ground_truth` file was written. `--graph` (which needs
        // `cluster_map`) is validated incompatible with this flag in
        // `main.rs`, so an empty map here is never actually consumed.
        None => crate::gt::GtResult {
            n_exact_dup: 0,
            n_fuzzy_dup: 0,
            n_hard_neg: 0,
            n_unique: 0,
            n_masters: 0,
            cluster_map: crate::gt::ClusterCsr::build(Vec::new()),
        },
    };
    let _t_gt_compute = 0.0f64;
    let _t_gt_write = t_gt0.elapsed().as_secs_f64();
    let t3b_elapsed = t3b.elapsed().as_secs_f64();
    let _t3_elapsed = t3.elapsed().as_secs_f64();
    log_rss("after GT finish (cluster_map built)");

    // ── Graph: emit duplicate-cluster edges from the post-GT cluster_map ──
    if config.graph_enabled
        && let Some(ew) = edge_writer.as_mut()
    {
        crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
            .map_err(|e| format!("push dup clusters: {e}"))?;
    }

    let duration = t_start.elapsed().as_secs_f64();

    log::debug!(
        "[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",
        t_alloc,
        t1e_elapsed,
        t2_elapsed,
        t3a_elapsed,
        _t_meta,
        _t_dup,
        _t_write,
        _t_gt_compute,
        _t_gt_write,
        t3b_elapsed,
        duration,
    );

    let stats = PipelineStats {
        total_records: global_rid_offset,
        exact_dups: n_exact_dup,
        fuzzy_dups: n_fuzzy_dup,
        hard_negs: n_hard_neg,
        uniques: n_unique,
        masters: n_masters,
    };

    // ── Finalize graph writers (base nodes + FK edges emitted so far) ──
    let mut graph_nodes_final: Option<String> = None;
    let mut graph_edges_final: Option<String> = None;
    if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
        nw.finish()
            .map_err(|e| format!("finish node writer: {e}"))?;
        ew.finish()
            .map_err(|e| format!("finish edge writer: {e}"))?;
        let mut np = graph_nodes_path.clone().unwrap();
        let mut ep = graph_edges_path.clone().unwrap();
        if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
            let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
            let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
            convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
            convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
            std::fs::remove_file(&np).ok();
            std::fs::remove_file(&ep).ok();
            np = np_pq;
            ep = ep_pq;
        }
        graph_nodes_final = Some(np);
        graph_edges_final = Some(ep);
    }
    log_rss("end (after graph finalize)");

    // Carry each entity's final local master_id counter forward for the
    // next chunk (`--chunk-size`) — pool_only entities never assign a
    // master_id (see the `continue` above), so their slot is left as-is.
    let mut master_local = offsets.master_local.clone();
    if master_local.len() < config.entity_plans.len() {
        master_local.resize(config.entity_plans.len(), 0);
    }
    for (idx, plan) in config.entity_plans.iter().enumerate() {
        if !plan.pool_only {
            master_local[idx] = offsets.master_base(idx) + plan.n_base;
        }
    }
    // Pool-only entities write no canary rows (see `canary::generate_all`).
    let id_local = config
        .entity_plans
        .iter()
        .enumerate()
        .map(|(idx, plan)| {
            let canaries = if plan.pool_only {
                0
            } else {
                crate::canary::CANARY_COUNT
            };
            offsets.id_base(idx) + plan.n_base + canaries
        })
        .collect();
    let next_offsets = ChunkOffsets {
        rid_offset: global_rid_offset,
        master_local,
        hn_master_counter: hn_master_id_counter,
        id_local,
    };

    Ok((
        PipelineOutput {
            output_files: vec![dataset_path],
            // Empty when `--skip-ground-truth` — no `_ground_truth` file was
            // ever written, so `gt_path` would point to nothing.
            gt_file: if config.skip_ground_truth {
                String::new()
            } else {
                gt_path
            },
            stats,
            nodes: graph_nodes_final,
            edges: graph_edges_final,
        },
        next_offsets,
    ))
}

/// IPC → Parquet (ZSTD) conversion for graph files, mirroring the dataset
/// conversion above. The node IPC already carries the `dupehell.*` metadata,
/// which is forwarded to the Parquet key/value metadata.
fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
    use arrow::ipc::reader::FileReader;

    let ipc_file =
        std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
    let ipc_reader =
        FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
    let schema = ipc_reader.schema();

    let parquet_file =
        std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
    let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
    // Sorted: `schema.metadata()` is a `HashMap` (see `RunMetadata`).
    let meta_kv = parquet_kv(&schema.metadata().clone().into_iter().collect());
    let props = parquet::file::properties::WriterProperties::builder()
        .set_compression(parquet::basic::Compression::ZSTD(zstd_level))
        .set_max_row_group_row_count(Some(1_000_000))
        .set_key_value_metadata(Some(meta_kv))
        .build();
    let mut parquet_writer =
        parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
            .map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;

    for batch_result in ipc_reader {
        let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
        parquet_writer
            .write(&batch)
            .map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
    }
    parquet_writer
        .close()
        .map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
    Ok(())
}

// ── Chunked generation (`--chunk-size`) ─────────────────────────────────────

/// Streams every batch (IPC) or row group (Parquet) of each file in
/// `in_paths` — all sharing the same schema, as chunk outputs from
/// `run_chunked` always do — into a single `out_path`, never holding more
/// than one batch/row-group resident at a time. Used to assemble
/// `--chunk-size` chunk outputs (dataset/GT/graph files) into the one final
/// file a normal (non-chunked) run would have produced.
fn concat_dataset_files(
    output_format: &str,
    in_paths: &[String],
    out_path: &str,
) -> Result<(), String> {
    if in_paths.is_empty() {
        return Err("concat_dataset_files: no chunk files to concatenate".into());
    }
    if output_format == "parquet" {
        concat_parquet_files(in_paths, out_path)
    } else {
        concat_ipc_files(in_paths, out_path)
    }
}

fn concat_ipc_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
    use arrow::ipc::reader::FileReader;

    let first_file =
        std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
    let schema = FileReader::try_new(first_file, None)
        .map_err(|e| format!("ipc reader {}: {e}", in_paths[0]))?
        .schema();

    let out_file =
        std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
    let mut writer = arrow::ipc::writer::FileWriter::try_new(out_file, &schema)
        .map_err(|e| format!("ipc writer {out_path}: {e}"))?;

    for path in in_paths {
        let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
        let reader =
            FileReader::try_new(file, None).map_err(|e| format!("ipc reader {path}: {e}"))?;
        for batch_result in reader {
            let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
            writer
                .write(&batch)
                .map_err(|e| format!("write {out_path}: {e}"))?;
        }
    }
    writer
        .finish()
        .map_err(|e| format!("finish {out_path}: {e}"))
}

fn concat_parquet_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

    let first_file =
        std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
    let first_builder = ParquetRecordBatchReaderBuilder::try_new(first_file)
        .map_err(|e| format!("parquet reader {}: {e}", in_paths[0]))?;
    let schema = first_builder.schema().clone();
    let meta_kv: Vec<parquet::file::metadata::KeyValue> = first_builder
        .metadata()
        .file_metadata()
        .key_value_metadata()
        .cloned()
        .unwrap_or_default();

    let out_file =
        std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
    let zstd = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
    let props = parquet::file::properties::WriterProperties::builder()
        .set_compression(parquet::basic::Compression::ZSTD(zstd))
        .set_max_row_group_row_count(Some(1_000_000))
        .set_key_value_metadata(Some(meta_kv))
        .build();
    let mut writer = parquet::arrow::ArrowWriter::try_new(out_file, schema, Some(props))
        .map_err(|e| format!("arrow writer {out_path}: {e}"))?;

    for path in in_paths {
        let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
        let reader = ParquetRecordBatchReaderBuilder::try_new(file)
            .map_err(|e| format!("parquet reader {path}: {e}"))?
            .build()
            .map_err(|e| format!("build parquet reader {path}: {e}"))?;
        for batch_result in reader {
            let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
            writer
                .write(&batch)
                .map_err(|e| format!("write {out_path}: {e}"))?;
        }
    }
    writer
        .close()
        .map(|_| ())
        .map_err(|e| format!("close {out_path}: {e}"))
}

/// Generates `total_size` records by internally looping over
/// sequential, independently-seeded chunk runs (see [`chunk_sizes`]),
/// keeping `record_id`/`master_id` globally
/// contiguous across chunks (see `ChunkOffsets`) instead of the shard-tag
/// hack an external concat script would need, then assembling the chunk
/// outputs into one final dataset/GT/graph file via `concat_dataset_files`.
/// RAM stays bounded to roughly one `--size chunk_size` run at a time — each
/// chunk is a fully independent pipeline run, only the identifier counters
/// and (for the final assembly step) the output files carry over.
///
/// If `chunk_size == 0` or `chunk_size >= total_size`, this is exactly one
/// chunk — delegates straight to `run_pipeline_with_progress`, byte-for-byte
/// identical to calling it directly (no `--chunk-size` given).
/// Smallest `--size` a single pipeline run accepts (see
/// `schema::build_pipeline_config`) — and so the smallest chunk.
pub const MIN_RUN_SIZE: usize = 10;

/// Sizes of the chunks `run_chunked` splits `total` into: `chunk_size`
/// each, except that a trailing remainder smaller than [`MIN_RUN_SIZE`] is
/// folded into the last full chunk instead of becoming a chunk of its own
/// (which `build_pipeline_config` would reject, after every previous chunk
/// had already been generated). Requires `chunk_size >= MIN_RUN_SIZE`.
pub fn chunk_sizes(total: usize, chunk_size: usize) -> Vec<usize> {
    let mut sizes = vec![chunk_size; total / chunk_size];
    let rem = total % chunk_size;
    match sizes.last_mut() {
        Some(last) if rem < MIN_RUN_SIZE => *last += rem,
        _ if rem > 0 => sizes.push(rem),
        _ => {}
    }
    sizes
}

struct RemoveDirOnDrop<'a>(&'a str);

impl Drop for RemoveDirOnDrop<'_> {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(self.0);
    }
}

#[allow(clippy::too_many_arguments)]
pub fn run_chunked(
    ctx: &Context,
    domain: &str,
    total_size: usize,
    chunk_size: usize,
    seed: u64,
    difficulty: &str,
    hard_neg_ratio: f64,
    singleton_master_fraction: f64,
    schema: &crate::schema::DomainSchema,
    final_run_id: &str,
    output_format: &str,
    graph_enabled: bool,
    graph_format: &str,
    only_entity: Option<&str>,
    output_dir: &str,
    mut progress: Option<&mut dyn FnMut(usize, usize)>,
    skip_ground_truth: bool,
) -> Result<PipelineOutput, String> {
    if chunk_size == 0 || chunk_size >= total_size {
        let mut config = crate::schema::build_pipeline_config(
            domain,
            total_size,
            seed,
            difficulty,
            hard_neg_ratio,
            singleton_master_fraction,
            schema,
            final_run_id,
            output_format,
            graph_enabled,
            graph_format,
            only_entity,
        )?;
        config.skip_ground_truth = skip_ground_truth;
        return run_pipeline_with_progress(ctx, &config, output_dir, progress);
    }

    if chunk_size < MIN_RUN_SIZE {
        return Err(format!(
            "--chunk-size must be >= {MIN_RUN_SIZE} (each chunk is a full run), got {chunk_size}"
        ));
    }
    let sizes = chunk_sizes(total_size, chunk_size);
    let tmp_dir = format!("{output_dir}/.dupehell_chunks_{final_run_id}");
    std::fs::create_dir_all(&tmp_dir).map_err(|e| format!("create {tmp_dir:?}: {e}"))?;
    // Removes the per-chunk files on every exit path, including a chunk
    // failing midway — not only after a successful assembly.
    let _tmp_guard = RemoveDirOnDrop(&tmp_dir);

    let mut offsets = ChunkOffsets::default();
    let mut dataset_chunks: Vec<String> = Vec::new();
    let mut gt_chunks: Vec<String> = Vec::new();
    let mut node_chunks: Vec<String> = Vec::new();
    let mut edge_chunks: Vec<String> = Vec::new();
    let mut agg = PipelineStats {
        total_records: 0,
        exact_dups: 0,
        fuzzy_dups: 0,
        hard_negs: 0,
        uniques: 0,
        masters: 0,
    };

    for (chunk_idx, &this_size) in sizes.iter().enumerate() {
        // Not `seed + chunk_idx`: that made chunk 1 of seed 42 the same
        // data as chunk 0 of seed 43.
        let chunk_seed =
            crate::rng::derive_seed(seed, &[crate::rng::stream::CHUNK, chunk_idx as u64]);
        let chunk_run_id = format!("{final_run_id}__chunk{chunk_idx:04}");
        let mut chunk_config = crate::schema::build_pipeline_config(
            domain,
            this_size,
            chunk_seed,
            difficulty,
            hard_neg_ratio,
            singleton_master_fraction,
            schema,
            &chunk_run_id,
            output_format,
            graph_enabled,
            graph_format,
            only_entity,
        )?;
        chunk_config.skip_ground_truth = skip_ground_truth;

        let already_written = agg.total_records;
        let mut chunk_cb = |done: usize, _total: usize| {
            if let Some(cb) = &mut progress {
                cb(already_written + done, total_size);
            }
        };
        let (out, next_offsets) = run_pipeline_chunked(
            ctx,
            &chunk_config,
            &tmp_dir,
            Some(&mut chunk_cb),
            Some(&offsets),
        )?;
        offsets = next_offsets;

        // `out.stats.total_records` is `global_rid_offset` at the end of
        // the chunk, which started from `offsets.rid_offset` (the previous
        // chunk's running total) — it's already the running grand total,
        // not a per-chunk delta.
        agg.total_records = out.stats.total_records;
        agg.exact_dups += out.stats.exact_dups;
        agg.fuzzy_dups += out.stats.fuzzy_dups;
        agg.hard_negs += out.stats.hard_negs;
        agg.uniques += out.stats.uniques;
        agg.masters += out.stats.masters;

        dataset_chunks.push(out.output_files[0].clone());
        if !skip_ground_truth {
            gt_chunks.push(out.gt_file.clone());
        }
        if let Some(n) = out.nodes {
            node_chunks.push(n);
        }
        if let Some(e) = out.edges {
            edge_chunks.push(e);
        }
    }

    let dataset_ext = if output_format == "parquet" {
        "parquet"
    } else {
        "ipc"
    };
    let final_dataset = format!("{output_dir}/{final_run_id}.{dataset_ext}");
    concat_dataset_files(output_format, &dataset_chunks, &final_dataset)?;

    let final_gt = if skip_ground_truth {
        String::new()
    } else {
        let p = format!("{output_dir}/{final_run_id}_ground_truth.{dataset_ext}");
        concat_dataset_files(output_format, &gt_chunks, &p)?;
        p
    };

    let graph_ext = if graph_format == "parquet" {
        "parquet"
    } else {
        "ipc"
    };
    let final_nodes = if !node_chunks.is_empty() {
        let p = format!("{output_dir}/{final_run_id}_nodes.{graph_ext}");
        concat_dataset_files(graph_format, &node_chunks, &p)?;
        Some(p)
    } else {
        None
    };
    let final_edges = if !edge_chunks.is_empty() {
        let p = format!("{output_dir}/{final_run_id}_edges.{graph_ext}");
        concat_dataset_files(graph_format, &edge_chunks, &p)?;
        Some(p)
    } else {
        None
    };

    Ok(PipelineOutput {
        output_files: vec![final_dataset],
        gt_file: final_gt,
        stats: agg,
        nodes: final_nodes,
        edges: final_edges,
    })
}

// ── Capacity limits and output-size estimate ────────────────────────────────

/// Digits of the `HN-` hard-negative `master_id` counter (see
/// `hn_master_id_array`): a run can number at most 10^9 hard negatives.
pub(crate) const HN_ID_DIGITS: usize = 9;

/// Rows this plan writes: base rows plus their duplicate copies.
fn plan_rows(plan: &EntityPlan) -> u64 {
    plan.n_base as u64 + plan.noise_types.iter().map(|n| n.count as u64).sum::<u64>()
}

/// Upper bound on the rows `--graph` keeps in the ground-truth cluster map
/// (every written row of a duplicated master): all base + duplicate rows.
/// Each costs ~24 bytes of RAM while the map is built (`gt::ClusterCsr`).
pub fn max_cluster_rows(config: &PipelineConfig) -> u64 {
    config
        .entity_plans
        .iter()
        .filter(|p| !p.pool_only)
        .map(plan_rows)
        .sum()
}

/// The output format's hard structural limits, checked before a run
/// instead of discovering them as silent corruption mid-run: fixed-width
/// identifiers that would wrap around (duplicate ids) and `--graph`'s
/// 32-bit cluster offsets. Generation itself no longer needs RAM in
/// proportion to `--size` (ground truth is streamed per batch), so these —
/// plus disk space, see [`estimate_output_bytes`] — are what bounds a run.
pub fn check_capacity(config: &PipelineConfig) -> Result<(), String> {
    let entity_id_limit = 10u64.pow(crate::column_gen::ENTITY_ID_DIGITS as u32);
    for plan in config.entity_plans.iter().filter(|p| !p.pool_only) {
        // + canary rows, numbered right after the base rows.
        if plan.n_base as u64 + 3 >= entity_id_limit {
            return Err(format!(
                "entity '{}' would get {} base rows, more than its {}-digit `_id` \
                 values can number uniquely ({entity_id_limit}) — use a smaller --size",
                plan.name,
                plan.n_base,
                crate::column_gen::ENTITY_ID_DIGITS
            ));
        }
    }
    let hn_limit = 10u64.pow(HN_ID_DIGITS as u32);
    let n_hn: u64 = config.hard_neg_types.iter().map(|h| h.count as u64).sum();
    if n_hn >= hn_limit {
        return Err(format!(
            "this run would generate {n_hn} hard negatives, more than their {HN_ID_DIGITS}-digit \
             master_id can number uniquely ({hn_limit}) — use a smaller --size or --hard-neg-ratio"
        ));
    }
    if config.graph_enabled {
        // Duplicate-cluster rows are indexed by `u32` offsets in
        // `gt::ClusterCsr`.
        let rows = max_cluster_rows(config);
        if rows > u32::MAX as u64 {
            return Err(format!(
                "--graph supports at most {} duplicate-cluster rows; this run could have up \
                 to {rows} — use a smaller --size, or drop --graph",
                u32::MAX
            ));
        }
    }
    Ok(())
}

// Per-row byte costs of an IPC dataset row, fitted (least squares) on
// measured runs of 7 schema shapes — aviation, aviation `--only-entity
// passenger`, ecommerce, fintech, healthcare, kyc, nonprofit (hell, 200K
// base rows): every estimate within ±2% of the real file size. A row costs
// a filled or a null string depending on whether the column belongs to its
// own entity (the dataset is the union of every entity's columns).
const IPC_FILLED_STR: f64 = 10.8;
const IPC_NULL_STR: f64 = 5.4;
const IPC_FIXED_WIDTH: f64 = 8.0;
const IPC_DICT: f64 = 4.0;
const IPC_BOOL: f64 = 0.25;
const IPC_ROW_CONST: f64 = 66.5;
// Same runs: Parquet (ZSTD(3)) datasets weighed 0.078-0.14x their IPC
// size; the upper end is used.
const PARQUET_RATIO: f64 = 0.15;
// Ground truth: 55.6 B/row (IPC), 3.1-3.4 B/row (Parquet), any domain.
const GT_IPC_ROW: f64 = 56.0;
const GT_PARQUET_ROW: f64 = 3.5;
// `--graph` edges: 87 B per dataset row in IPC on fintech (5 entities,
// FK-dense) — rounded up.
const EDGES_IPC_ROW: f64 = 90.0;

/// Estimated peak disk usage of a run, in bytes: dataset + ground truth +
/// graph (whose IPC files exist even in Parquet mode, until converted),
/// doubled for `--chunk-size` (chunk files stay on disk until assembled).
/// Calibrated on real runs (see the constants above); used for the
/// up-front free-space check, backed by `disk::DiskGuard` during the run.
pub fn estimate_output_bytes(config: &PipelineConfig, chunked: bool) -> u64 {
    let schema = build_full_schema(config, &RunMetadata::new());
    let fields = &schema.fields()[4..]; // after record_id/domain/entity_type/master_id
    let mut rows_by_entity: HashMap<&str, f64> = config
        .entity_plans
        .iter()
        .filter(|p| !p.pool_only)
        .map(|p| (p.name.as_str(), plan_rows(p) as f64))
        .collect();
    for hn in &config.hard_neg_types {
        if let Some(r) = rows_by_entity.get_mut(hn.entity_type.as_str()) {
            *r += hn.count as f64;
        }
    }

    let mut dataset_ipc = 0.0;
    let mut total_rows = 0.0;
    for plan in config.entity_plans.iter().filter(|p| !p.pool_only) {
        let rows = rows_by_entity[plan.name.as_str()];
        let own: std::collections::HashSet<String> =
            serde_json::from_str::<Vec<serde_json::Value>>(&plan.columns_json)
                .unwrap_or_default()
                .iter()
                .filter_map(|c| c["name"].as_str().map(str::to_string))
                .collect();
        let row_bytes: f64 = IPC_ROW_CONST
            + fields
                .iter()
                .map(|f| match f.data_type() {
                    DataType::Utf8 if own.contains(f.name()) => IPC_FILLED_STR,
                    DataType::Utf8 => IPC_NULL_STR,
                    DataType::Dictionary(_, _) => IPC_DICT,
                    DataType::Boolean => IPC_BOOL,
                    _ => IPC_FIXED_WIDTH,
                })
                .sum::<f64>();
        dataset_ipc += rows * row_bytes;
        total_rows += rows;
    }

    let parquet = config.output_format == "parquet";
    let mut bytes = if parquet {
        dataset_ipc * PARQUET_RATIO
    } else {
        dataset_ipc
    };
    if !config.skip_ground_truth {
        bytes += total_rows * if parquet { GT_PARQUET_ROW } else { GT_IPC_ROW };
    }
    if config.graph_enabled {
        let graph_ipc = dataset_ipc + total_rows * EDGES_IPC_ROW;
        bytes += graph_ipc;
        if config.graph_format == "parquet" {
            bytes += graph_ipc * PARQUET_RATIO;
        }
    }
    if chunked {
        bytes *= 2.0;
    }
    bytes as u64
}

// ── Metadata injection ──────────────────────────────────────────────────────

/// Run-level `dupehell.*` metadata embedded in every output file. A
/// `BTreeMap`, not a `HashMap`: writers that copy it entry by entry (Parquet
/// key/value metadata) then emit it in a stable, sorted key order — a
/// `HashMap`'s iteration order is randomized per process, which made two
/// otherwise identical Parquet files differ byte-wise.
pub(crate) type RunMetadata = std::collections::BTreeMap<String, String>;

/// `RunMetadata` as the `HashMap` Arrow's `Schema::with_metadata` takes
/// (Arrow's IPC writer sorts the keys itself on serialization).
pub(crate) fn arrow_metadata(metadata: &RunMetadata) -> HashMap<String, String> {
    metadata
        .iter()
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect()
}

/// Parquet key/value metadata, in sorted key order.
pub(crate) fn parquet_kv(metadata: &RunMetadata) -> Vec<parquet::file::metadata::KeyValue> {
    metadata
        .iter()
        .map(|(k, v)| parquet::file::metadata::KeyValue {
            key: k.clone(),
            value: Some(v.clone()),
        })
        .collect()
}

/// `SOURCE_DATE_EPOCH` (the reproducible-builds convention): when set to a
/// Unix timestamp, `dupehell.timestamp` uses it instead of the wall clock,
/// so two identical runs produce byte-identical files — that timestamp is
/// the only part of the output that otherwise depends on when it was
/// generated. `Ok(None)` when unset, `Err` when set to something that isn't
/// a non-negative integer (the CLI rejects that up front rather than
/// silently falling back to the clock).
pub fn source_date_epoch() -> Result<Option<u64>, String> {
    match std::env::var("SOURCE_DATE_EPOCH") {
        Err(_) => Ok(None),
        Ok(s) => s.trim().parse::<u64>().map(Some).map_err(|_| {
            format!("SOURCE_DATE_EPOCH must be a Unix timestamp (non-negative integer), got {s:?}")
        }),
    }
}

fn build_metadata_map(config: &PipelineConfig) -> RunMetadata {
    let ts = match source_date_epoch() {
        Ok(Some(epoch)) => epoch.to_string(),
        _ => std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs().to_string())
            .unwrap_or_else(|_| "unknown".into()),
    };
    RunMetadata::from([
        (
            "dupehell.generator".into(),
            format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
        ),
        (
            "dupehell.provenance".into(),
            "dupehell-synthetic-data".into(),
        ),
        ("dupehell.license".into(), "MIT".into()),
        (
            "dupehell.purpose".into(),
            "Educational Use Only — Record Linkage Benchmarking".into(),
        ),
        (
            "dupehell.url".into(),
            "https://github.com/vntoinekaio/DupeHell".into(),
        ),
        ("dupehell.domain".into(), config.domain.clone()),
        ("dupehell.size".into(), config.size.to_string()),
        ("dupehell.seed".into(), config.seed.to_string()),
        ("dupehell.difficulty".into(), config.difficulty.clone()),
        ("dupehell.run_id".into(), config.run_id.clone()),
        ("dupehell.timestamp".into(), ts),
        (
            "dupehell.structural_key_columns".into(),
            structural_key_columns_note(config),
        ),
    ])
}

/// Lists the `{entity}_id` columns that are structural join keys (linking an
/// entity's rows to its flattened child tables), not attributes to feed an
/// ER model — they never receive noise and stay identical across all
/// duplicates of the same master_id, unlike `record_id`.
fn structural_key_columns_note(config: &PipelineConfig) -> String {
    let cols: Vec<&str> = config
        .entity_plans
        .iter()
        .filter_map(|p| p.identifier_col.as_deref())
        .collect();
    format!(
        "{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
        cols.join(", ")
    )
}

// ── Schema alignment ───────────────────────────────────────────────────────

/// Build the union schema from all entity plans (all columns + metadata).
/// `Dictionary(Int32, Utf8)` for the near-constant-cardinality metadata
/// columns (`domain`/`entity_type`/`match_type`/`difficulty`/`edge_type`/
/// `subtype`) — an explicit output-contract change (not bit-identical:
/// the declared column type changes from `Utf8`).
pub(crate) fn low_cardinality_dict_type() -> DataType {
    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
}

/// A single dictionary (`values` + a `value -> key` index) shared by every
/// `DictionaryArray` built for one column across an entire output file.
///
/// Arrow's IPC `FileWriter` requires exactly one dictionary per field for
/// the whole file — writing a batch whose dictionary *content* differs from
/// a previously-written one for that field is a hard error ("Dictionary
/// replacement detected"), not just wasted work. A fresh
/// `StringDictionaryBuilder` per batch (the naive approach) builds a new,
/// batch-scoped dictionary every time — even when every batch happens to
/// use the exact same logical value set, the *order*/*subset* actually
/// registered can differ batch to batch, which is a real content
/// difference. `DictValues` fixes the dictionary's content once (built
/// from the full known value set for this column *before* any batch is
/// written) and every batch's `DictionaryArray` reuses the same `values`
/// `Arc` — trivially identical content by construction, for the life of
/// one output file.
pub(crate) struct DictValues {
    values: ArrayRef,
    index: HashMap<String, i32>,
}

impl DictValues {
    pub(crate) fn new<I, S>(values: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let owned: Vec<String> = values.into_iter().map(Into::into).collect();
        let index: HashMap<String, i32> = owned
            .iter()
            .enumerate()
            .map(|(i, v)| (v.clone(), i as i32))
            .collect();
        Self {
            values: Arc::new(StringArray::from(owned)),
            index,
        }
    }

    fn idx(&self, value: &str) -> i32 {
        *self.index.get(value).unwrap_or_else(|| {
            panic!("DictValues: value {value:?} was not in the pre-built dictionary")
        })
    }

    fn build(&self, keys: arrow::array::Int32Array) -> ArrayRef {
        Arc::new(
            arrow::array::DictionaryArray::<arrow::datatypes::Int32Type>::try_new(
                keys,
                self.values.clone(),
            )
            .expect("DictValues::build: key array out of range for this dictionary"),
        )
    }

    /// `n` copies of `value` (must already be in the dictionary).
    pub(crate) fn const_array(&self, value: &str, n: usize) -> ArrayRef {
        self.build(arrow::array::Int32Array::from(vec![self.idx(value); n]))
    }

    /// Exposed for callers that build a `DictionaryArray` row by row (e.g.
    /// `gt.rs`'s `match_type` column): the key for `value`, to push into an
    /// `Int32Builder`, plus `finish_keys` to wrap the finished key array in
    /// this dictionary's `values`.
    pub(crate) fn key(&self, value: &str) -> i32 {
        self.idx(value)
    }

    pub(crate) fn finish_keys(&self, keys: arrow::array::Int32Array) -> ArrayRef {
        self.build(keys)
    }

    /// `n` all-null entries sharing this dictionary's `values`: used when a batch's own schema is missing a
    /// dictionary-encoded column entirely (e.g. an HN pattern whose config
    /// never references it) — `arrow::array::new_null_array` would build a
    /// *different*, freshly-constructed dictionary values array for that
    /// batch, which IPC then rejects as "Dictionary replacement detected"
    /// against every other batch of the same field that uses the real
    /// shared dictionary. All-null keys map to no value regardless of which
    /// `values` array they're paired with, so sharing this one costs
    /// nothing and keeps every batch of the field on the same dictionary.
    pub(crate) fn null_array(&self, n: usize) -> ArrayRef {
        self.build(arrow::array::Int32Array::new_null(n))
    }
}

/// Full known `subtype` value set for `_edges` output: every
/// FK remap's `source_col` (the `"fk"`-edge subtype, see the `push()` call
/// in the base-batch FK-edge loop above), every hard-negative pattern kind
/// actually configured (the `"hard_neg"`-edge subtype — `config_json`'s
/// `"pattern"` field, e.g. `"same_field"`, NOT the schema's per-type name;
/// see `hn_common::generate_hard_negatives`), and the two literal dup-
/// cluster edge subtypes (`graph_gen::push_dup_clusters`). Built once,
/// before any edge is written, so `EdgeWriter` can hold one dictionary for
/// the whole `_edges` file's lifetime.
fn subtype_dict(config: &PipelineConfig) -> DictValues {
    let mut values: Vec<String> = vec!["complete".to_string(), "spanning_tree".to_string()];
    for plan in &config.entity_plans {
        for remap in &plan.fk_remaps {
            values.push(remap.source_col.clone());
        }
    }
    for hn in &config.hard_neg_types {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&hn.config_json)
            && let Some(pattern) = v.get("pattern").and_then(|p| p.as_str())
        {
            values.push(pattern.to_string());
        }
    }
    values.sort_unstable();
    values.dedup();
    DictValues::new(values)
}

fn build_full_schema(config: &PipelineConfig, metadata: &RunMetadata) -> Schema {
    let mut field_map: Vec<(String, DataType, bool)> = vec![
        ("record_id".to_string(), DataType::Utf8, false),
        ("domain".to_string(), low_cardinality_dict_type(), false),
        (
            "entity_type".to_string(),
            low_cardinality_dict_type(),
            false,
        ),
        ("master_id".to_string(), DataType::Utf8, false),
    ];
    for plan in &config.entity_plans {
        // Pool-only entities (`--only-entity`) are never written to output —
        // excluding their columns here is what gives `--only-entity` its
        // narrow, single-entity schema for free, with no separate filtering
        // step downstream.
        if plan.pool_only {
            continue;
        }
        let cols: Vec<serde_json::Value> =
            serde_json::from_str(&plan.columns_json).unwrap_or_default();
        for col in &cols {
            let name = col["name"].as_str().unwrap_or("").to_string();
            if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
                continue;
            }
            // The declared schema type must match what `generate_column`
            // actually builds, or `RecordBatch::
            // try_new` below hard-errors on a type mismatch. A handful of
            // templates produce `Dictionary(Int32, Utf8)` instead of plain
            // `Utf8` — `resolve_dict_encoded_column` is the single place
            // that name resolution is written, shared with
            // `difficulty::estimate_difficulty`'s own model of the same
            // fact.
            let is_dict = crate::fast_template::resolve_dict_encoded_column(&name).is_some();
            let col_type = if is_dict {
                low_cardinality_dict_type()
            } else {
                col_type_from_request(col)
            };
            field_map.push((name, col_type, true));
        }
    }
    let fields: Vec<Field> = field_map
        .into_iter()
        .map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
        .collect();
    Schema::new(fields).with_metadata(arrow_metadata(metadata))
}

/// Map a column's JSON type string to Arrow DataType.
fn col_type_from_request(v: &serde_json::Value) -> DataType {
    match v["type"].as_str() {
        Some("int") => DataType::Int64,
        Some("float") => DataType::Float64,
        Some("boolean") => DataType::Boolean,
        Some("date") | Some("datetime") => DataType::Utf8,
        _ => DataType::Utf8,
    }
}

// ── Helpers ─────────────────────────────────────────────────────────────────

#[allow(clippy::too_many_arguments)]
pub(crate) fn add_metadata_and_align(
    rb: &RecordBatch,
    domain: &str,
    entity_type: &str,
    domain_dict: &DictValues,
    entity_type_dict: &DictValues,
    record_ids: &ArrayRef,
    master_ids: &ArrayRef,
    full_arc: &Arc<Schema>,
    col_lookup: &[Option<usize>],
    null_cache: &mut FxHashMap<(DataType, usize), ArrayRef>,
    const_arr_cache: &mut FxHashMap<(String, usize), ArrayRef>,
) -> Result<RecordBatch, String> {
    let n = rb.num_rows();
    // `domain` is constant for the whole run and `entity_type` repeats
    // across every batch of a given entity (and often across HN/canary
    // sections too) — cache the built repeat-array per (value, n) instead
    // of rebuilding an n-length array on every call, same idea as
    // `null_cache` above for the always-null columns. Built as
    // `Dictionary(Int32, Utf8)` against the run-wide `domain_dict`/
    // `entity_type_dict` — every batch's array reuses the
    // exact same dictionary `values`, required for IPC output (see
    // `DictValues`'s doc comment).
    let domain_arr = const_arr_cache
        .entry((domain.to_string(), n))
        .or_insert_with(|| domain_dict.const_array(domain, n))
        .clone();
    let et_arr = const_arr_cache
        .entry((entity_type.to_string(), n))
        .or_insert_with(|| entity_type_dict.const_array(entity_type, n))
        .clone();
    let rid_arr = record_ids.clone();
    let mid_arr = master_ids.clone();

    let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
    all_arrays.push(rid_arr);
    all_arrays.push(domain_arr);
    all_arrays.push(et_arr);
    all_arrays.push(mid_arr);

    // P4: Vec<Option<usize>> eliminates HashMap lookup per field (just index)
    // Phase 12.3: cache null arrays per (DataType, n) — avoids new_null_array overhead per batch
    for (i, maybe_idx) in col_lookup.iter().enumerate() {
        match maybe_idx {
            Some(idx) => all_arrays.push(rb.column(*idx).clone()),
            None => {
                let field = full_arc.field(i + 4);
                // A dictionary-encoded column needs its
                // dedicated shared dictionary even when null, not the
                // general `(DataType, n)` cache below: two different
                // dict-encoded fields (e.g. `suffix` and `grade`) share the
                // exact same `DataType::Dictionary(Int32, Utf8)`, so
                // caching by `DataType` alone would hand one field's null
                // array to the other — wrong dictionary content, same
                // "Dictionary replacement" IPC error this fallback exists
                // to avoid in the first place.
                if let Some(arr) = crate::fast_template::dict_encoded_null_array(field.name(), n) {
                    all_arrays.push(arr);
                    continue;
                }
                let dt = field.data_type();
                let arr = null_cache
                    .entry((dt.clone(), n))
                    .or_insert_with(|| arrow::array::new_null_array(dt, n));
                all_arrays.push(Arc::clone(arr));
            }
        }
    }

    // P3: Arc::clone(full_arc) avoids full schema clone (just atomic increment)
    RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
        .map_err(|e| format!("add_metadata_and_align RecordBatch: {e}"))
}

fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
    use arrow::compute::take;
    let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
        .map(|i| take(rb.column(i), indices, None))
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| format!("take error: {e}"))?;

    RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::{build_pipeline_config, load_schema};

    /// Diagnostic scan for noise coverage gaps: for every entity of every domain schema,
    /// checks whether AT LEAST ONE of `hell`'s 9 noise categories has any
    /// target column at all. An entity with zero matching columns across
    /// all 9 categories gets `target_cols.is_empty()` on every attempt
    /// (`apply_noise_with_retry`, line ~811) -- its duplicates stay
    /// byte-identical to their master forever, regardless of `passes`, on
    /// every difficulty tier (not just hell, since hell's category list is
    /// the broadest; an entity failing all 9 fails every subset too).
    /// Confirmed root cause for aviation's `flight` entity (0/9
    /// categories match `flight_number`/`origin_airport`/
    /// `destination_airport`/`departure_time`/`arrival_time`/`status`/
    /// `source_system` -- `departure_time`/`arrival_time` are `datetime`
    /// columns but the `dates` category predicate only matches
    /// `"date"`/`"birth"`/`"incorporation"`/`"founding"`, never `"time"`).
    ///
    /// Deliberately reimplements just enough of `col_type_from_str`'s
    /// Utf8-vs-not split (only int/float/boolean are excluded; string,
    /// date, and datetime all end up `DataType::Utf8` per
    /// `entity_gen::col_type_to_arrow`) rather than importing it, since
    /// that function is private to `entity_gen` -- diagnostic-only code,
    /// not meant to become a permanent dependency.
    ///
    /// This is a report, not an assertion: intentionally does NOT fail the
    /// suite (a coverage gap is a data-quality finding to investigate,
    /// not a regression this specific test should gate) -- run with
    /// `cargo test --release diag_noise_coverage_scan -- --nocapture --ignored`
    /// to see the full per-domain/per-entity table.
    #[test]
    #[ignore = "diagnostic scan, not a pass/fail regression test -- run explicitly with --ignored --nocapture"]
    fn diag_noise_coverage_scan() {
        const HELL_NOISE_TYPES: &[&str] = &[
            "typo_aggressive",
            "visual",
            "unicode_pollution",
            "names",
            "dates",
            "identifiers",
            "addresses",
            "companies",
            "extra",
        ];

        let mut zero_coverage: Vec<(String, String, usize)> = Vec::new();
        let mut total_entities = 0usize;
        let mut schema_files: Vec<_> = std::fs::read_dir("schemas")
            .expect("read schemas/ dir")
            .filter_map(|e| e.ok())
            .map(|e| e.path())
            .filter(|p| p.extension().is_some_and(|e| e == "json"))
            .collect();
        schema_files.sort();

        for path in &schema_files {
            let domain = path.file_stem().unwrap().to_string_lossy().to_string();
            let text = std::fs::read_to_string(path).unwrap_or_else(|e| panic!("{path:?}: {e}"));
            let json: serde_json::Value =
                serde_json::from_str(&text).unwrap_or_else(|e| panic!("{path:?}: {e}"));
            let Some(entities) = json.get("entities").and_then(|v| v.as_array()) else {
                continue;
            };
            for entity in entities {
                total_entities += 1;
                let ename = entity
                    .get("name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("?")
                    .to_string();
                let fk_exclude: std::collections::HashSet<&str> = entity
                    .get("fk_remaps")
                    .and_then(|v| v.as_array())
                    .map(|arr| {
                        arr.iter()
                            .filter_map(|r| r.get("source_col").and_then(|v| v.as_str()))
                            .collect()
                    })
                    .unwrap_or_default();
                let cols = entity
                    .get("columns")
                    .and_then(|v| v.as_array())
                    .cloned()
                    .unwrap_or_default();

                let mut n_matched_categories = 0usize;
                for &ntype in HELL_NOISE_TYPES {
                    let has_match = cols.iter().any(|c| {
                        let name = c.get("name").and_then(|v| v.as_str()).unwrap_or("");
                        let ctype = c.get("type").and_then(|v| v.as_str()).unwrap_or("string");
                        let is_utf8 = !matches!(ctype, "int" | "float" | "boolean");
                        is_utf8
                            && !fk_exclude.contains(name)
                            && noise_type_targets_column(ntype, name)
                    });
                    if has_match {
                        n_matched_categories += 1;
                    }
                }
                if n_matched_categories == 0 {
                    zero_coverage.push((domain.clone(), ename, cols.len()));
                }
            }
        }

        eprintln!(
            "\n[noise_coverage] {total_entities} entities scanned across {} domains",
            schema_files.len()
        );
        eprintln!(
            "[noise_coverage] {} entities with ZERO matching hell noise category:",
            zero_coverage.len()
        );
        for (domain, entity, ncols) in &zero_coverage {
            eprintln!("[noise_coverage]   {domain}::{entity} ({ncols} columns)");
        }
    }

    /// Regression for the noise-category cross-contamination bug:
    /// `COMPANY_NAME_WORDS` and `PERSON_NAME_WORDS` both contain the bare
    /// `"name"` fragment, which used to let `given_name`/`family_name`/
    /// `middle_name` match the `companies` noise category (and would
    /// symmetrically let `company_name`/`legal_name` match `names`), and let
    /// `email_address` match the `addresses` category via `"address"`.
    #[test]
    fn test_no_person_company_name_cross_contamination() {
        for col in [
            "given_name",
            "family_name",
            "middle_name",
            "first_name",
            "last_name",
        ] {
            assert!(
                noise_type_targets_column("names", col),
                "{col} should still match 'names'"
            );
            assert!(
                !noise_type_targets_column("companies", col),
                "{col} should NOT match 'companies' (person name, not a company name)"
            );
        }
        for col in ["company_name", "legal_name", "trading_name"] {
            assert!(
                noise_type_targets_column("companies", col),
                "{col} should still match 'companies'"
            );
            assert!(
                !noise_type_targets_column("names", col),
                "{col} should NOT match 'names' (company name, not a person name)"
            );
        }
        assert!(
            !noise_type_targets_column("addresses", "email_address"),
            "email_address should NOT match 'addresses' (it's an email, not a postal address)"
        );
        assert!(
            noise_type_targets_column("identifiers", "email_address"),
            "email_address should still match 'identifiers' (corrupt_email)"
        );
        assert!(
            noise_type_targets_column("addresses", "residential_address_line1"),
            "a real postal address column should still match 'addresses'"
        );
    }

    fn manifest_dir() -> std::path::PathBuf {
        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).to_path_buf()
    }

    /// Unique per-test scratch dir under `target/`, so parallel test runs
    /// never collide on output filenames and cleanup is trivial.
    fn scratch_dir(name: &str) -> std::path::PathBuf {
        let dir = manifest_dir()
            .join("target")
            .join("pipeline_test_tmp")
            .join(format!("{name}_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).expect("create scratch dir");
        dir
    }

    fn run_kyc(
        output_dir: &std::path::Path,
        output_format: &str,
        graph_enabled: bool,
    ) -> PipelineOutput {
        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
        let mut ctx = Context::new(
            "kyc",
            "en",
            &manifest_dir().join("assets/pools").to_string_lossy(),
        )
        .expect("load context");
        let config = build_pipeline_config(
            "kyc",
            200,
            42,
            "medium",
            0.1,
            0.3,
            &schema,
            "pipeline_test_run",
            output_format,
            graph_enabled,
            "parquet",
            None,
        )
        .expect("build config");
        ctx.enable_watermark(&config.domain, config.size, config.seed);
        run_pipeline(&ctx, &config, &output_dir.to_string_lossy()).expect("run_pipeline")
    }

    #[test]
    fn test_run_pipeline_produces_dataset_and_gt() {
        let dir = scratch_dir("basic");
        let output = run_kyc(&dir, "parquet", false);

        assert_eq!(output.output_files.len(), 1);
        for f in &output.output_files {
            assert!(std::path::Path::new(f).exists(), "missing dataset file {f}");
        }
        assert!(std::path::Path::new(&output.gt_file).exists());
        assert!(output.nodes.is_none());
        assert!(output.edges.is_none());

        // Stats must be internally consistent: exact_dup/fuzzy_dup/hard_neg/
        // unique are a partition of the classified records, which is <=
        // total_records (total_records also includes unclassified canary
        // watermark rows).
        let s = &output.stats;
        assert!(s.total_records >= s.exact_dups + s.fuzzy_dups + s.hard_negs + s.uniques);
        assert!(s.total_records > 0);
        assert!(s.masters > 0);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_run_pipeline_ipc_format() {
        let dir = scratch_dir("ipc");
        let output = run_kyc(&dir, "ipc", false);
        assert!(output.output_files[0].ends_with(".ipc"));
        assert!(std::path::Path::new(&output.output_files[0]).exists());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_run_pipeline_graph_output() {
        let dir = scratch_dir("graph");
        let output = run_kyc(&dir, "parquet", true);
        let nodes = output.nodes.expect("nodes path");
        let edges = output.edges.expect("edges path");
        assert!(std::path::Path::new(&nodes).exists());
        assert!(std::path::Path::new(&edges).exists());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_fk_edge_weight_from_schema_and_default() {
        // aviation's flight->aircraft fk_remap declares fk_weight=0.5
        // (schemas/aviation.json); flight->airline leaves the field at
        // its schema default of 1.0. Confirms `FkRemapEntry::fk_weight`
        // is threaded through to the `weight` column of `fk`-typed graph
        // edges instead of the old hardcoded 1.0 constant, and that
        // domains/relations without the field still default correctly.
        let dir = scratch_dir("fk_weight");
        let schema =
            load_schema("aviation", &manifest_dir().join("schemas")).expect("load aviation.json");
        let mut ctx = Context::new(
            "aviation",
            "en",
            &manifest_dir().join("assets/pools").to_string_lossy(),
        )
        .expect("load context");
        let config = build_pipeline_config(
            "aviation",
            300,
            42,
            "medium",
            0.1,
            0.3,
            &schema,
            "pipeline_test_fk_weight",
            "ipc",
            true,
            "ipc",
            None,
        )
        .expect("build config");
        ctx.enable_watermark(&config.domain, config.size, config.seed);
        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");

        let edges_path = output.edges.expect("edges path");
        let file = std::fs::File::open(&edges_path).expect("open edges file");
        let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");

        let mut subtypes_seen: std::collections::HashMap<String, f64> =
            std::collections::HashMap::new();
        for batch in reader {
            let batch = batch.expect("read edges batch");
            // `edge_type`/`subtype` are `Dictionary(Int32, Utf8)`
            // — cast back to plain `Utf8` for this test-only
            // comparison.
            let edge_type_col = arrow::compute::cast(
                batch.column_by_name("edge_type").expect("edge_type col"),
                &DataType::Utf8,
            )
            .expect("cast edge_type to Utf8");
            let edge_type = edge_type_col.as_string::<i32>();
            let subtype_col = arrow::compute::cast(
                batch.column_by_name("subtype").expect("subtype col"),
                &DataType::Utf8,
            )
            .expect("cast subtype to Utf8");
            let subtype = subtype_col.as_string::<i32>();
            let weight = batch
                .column_by_name("weight")
                .expect("weight col")
                .as_any()
                .downcast_ref::<arrow::array::Float64Array>()
                .expect("weight as f64");
            for i in 0..batch.num_rows() {
                if edge_type.value(i) == "fk" {
                    subtypes_seen.insert(subtype.value(i).to_string(), weight.value(i));
                }
            }
        }

        assert_eq!(
            subtypes_seen.get("aircraft_id"),
            Some(&0.5),
            "flight->aircraft fk_weight from schema not propagated to edge weight"
        );
        assert_eq!(
            subtypes_seen.get("airline_id"),
            Some(&1.0),
            "flight->airline fk_weight (schema default) not propagated to edge weight"
        );

        let _ = std::fs::remove_dir_all(&dir);
    }

    /// End-to-end regression for the `ClusterCsr` rewrite of `cluster_map`:
    /// drives the real pipeline
    /// (`gt::GtAccumulator::finish` -> `ClusterCsr::build` ->
    /// `graph_gen::push_dup_clusters`) instead of testing either module in
    /// isolation, and checks the emitted `exact_dup`/`fuzzy_dup` edges are
    /// well-formed and consistent with the run's own GT stats.
    #[test]
    fn test_run_pipeline_dup_cluster_edges_well_formed() {
        let dir = scratch_dir("dup_cluster_edges");
        // `run_kyc` always requests parquet-format graph output regardless
        // of its `output_format` arg, so build the config directly here to
        // get IPC edges (matches `test_fk_edge_weight_from_schema_and_default`'s
        // approach for the same reason).
        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
        let mut ctx = Context::new(
            "kyc",
            "en",
            &manifest_dir().join("assets/pools").to_string_lossy(),
        )
        .expect("load context");
        let config = build_pipeline_config(
            "kyc",
            200,
            42,
            "medium",
            0.1,
            0.3,
            &schema,
            "pipeline_test_dup_cluster_edges",
            "ipc",
            true,
            "ipc",
            None,
        )
        .expect("build config");
        ctx.enable_watermark(&config.domain, config.size, config.seed);
        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");

        let edges_path = output.edges.expect("edges path");
        let file = std::fs::File::open(&edges_path).expect("open edges file");
        let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");

        let mut n_exact = 0usize;
        let mut n_fuzzy = 0usize;
        for batch in reader {
            let batch = batch.expect("read edges batch");
            let src = batch
                .column_by_name("source_node_id")
                .expect("source_node_id col")
                .as_any()
                .downcast_ref::<arrow::array::StringArray>()
                .expect("source_node_id as string");
            let tgt = batch
                .column_by_name("target_node_id")
                .expect("target_node_id col")
                .as_any()
                .downcast_ref::<arrow::array::StringArray>()
                .expect("target_node_id as string");
            // `edge_type` is `Dictionary(Int32, Utf8)` —
            // cast back to plain `Utf8` for this test-only comparison.
            let edge_type_col = arrow::compute::cast(
                batch.column_by_name("edge_type").expect("edge_type col"),
                &DataType::Utf8,
            )
            .expect("cast edge_type to Utf8");
            let edge_type = edge_type_col.as_string::<i32>();
            let weight = batch
                .column_by_name("weight")
                .expect("weight col")
                .as_any()
                .downcast_ref::<arrow::array::Float64Array>()
                .expect("weight as f64");

            for i in 0..batch.num_rows() {
                match edge_type.value(i) {
                    "exact_dup" => n_exact += 1,
                    "fuzzy_dup" => n_fuzzy += 1,
                    _ => continue,
                }
                // Every endpoint must round-trip through `parse_record_idx`
                // -- i.e. `ClusterCsr`'s packed indices were correctly
                // unpacked back into real record_ids by `push_dup_clusters`,
                // not garbage or truncated.
                assert!(
                    parse_record_idx(src.value(i)).is_some(),
                    "malformed source_node_id {:?} on a dup-cluster edge",
                    src.value(i)
                );
                assert!(
                    parse_record_idx(tgt.value(i)).is_some(),
                    "malformed target_node_id {:?} on a dup-cluster edge",
                    tgt.value(i)
                );
                assert!((weight.value(i) - 1.0).abs() < 1e-9);
            }
        }

        // kyc/medium/seed 42/size 200 always produces some duplicates (see
        // `output.stats.exact_dups`/`fuzzy_dups`), so the cluster-edge pass
        // must actually have emitted something -- an empty result here would
        // mean `ClusterCsr` silently dropped every cluster.
        assert!(n_exact + n_fuzzy > 0, "no dup-cluster edges emitted");
        assert!(output.stats.exact_dups + output.stats.fuzzy_dups > 0);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_run_pipeline_deterministic_stats() {
        let dir_a = scratch_dir("det_a");
        let dir_b = scratch_dir("det_b");
        let a = run_kyc(&dir_a, "parquet", false);
        let b = run_kyc(&dir_b, "parquet", false);
        assert_eq!(a.stats.total_records, b.stats.total_records);
        assert_eq!(a.stats.exact_dups, b.stats.exact_dups);
        assert_eq!(a.stats.fuzzy_dups, b.stats.fuzzy_dups);
        assert_eq!(a.stats.hard_negs, b.stats.hard_negs);
        assert_eq!(a.stats.uniques, b.stats.uniques);
        assert_eq!(a.stats.masters, b.stats.masters);
        let _ = std::fs::remove_dir_all(&dir_a);
        let _ = std::fs::remove_dir_all(&dir_b);
    }

    #[test]
    fn test_distribute_by_weight_sums_to_total_and_is_proportional() {
        // Regression: duplicates used to be sampled only from the last
        // batch of masters (`last_batch`), so for any entity spanning 2+
        // batches, earlier batches never received a single duplicate —
        // the observed symptom was the overall duplicate rate silently
        // drifting down as dataset size grew past `BATCH_SIZE`. Fixed by
        // spreading each noise_type's count across every batch in
        // proportion to its size via this helper.
        let weights = [500_000usize, 216_666];
        let counts = distribute_by_weight(283_334, &weights);
        assert_eq!(counts.len(), 2);
        assert_eq!(counts.iter().sum::<usize>(), 283_334);
        // Every batch gets a non-zero share proportional to its size —
        // in particular the *first* batch (which the old code starved
        // entirely) must receive the majority share here.
        assert!(counts[0] > 0, "first batch got zero duplicates");
        assert!(counts[1] > 0, "last batch got zero duplicates");
        assert!(counts[0] > counts[1]);

        // Edge cases: zero total, zero weights.
        assert_eq!(distribute_by_weight(0, &weights), vec![0, 0]);
        assert_eq!(distribute_by_weight(100, &[0, 0]), vec![0, 0]);
    }

    #[test]
    fn test_entity_prefix_unique_per_index() {
        let a = entity_prefix(0);
        let b = entity_prefix(1);
        assert_ne!(a, b);
        assert_eq!(a, "E00000");
        assert_eq!(entity_prefix(42), "E00042");
    }

    #[test]
    fn test_run_pipeline_hell_multi_pass_touches_multiple_columns() {
        // At `hell` (3 noise passes, see `schema::DifficultySettings::
        // passes`), at least one duplicate copy should differ from its
        // master on 2+ distinct columns — a single pass only ever touches
        // the columns its one chosen noise_type's predicate matches
        // (`noise_type_targets_column`), so seeing 2+ changed columns on
        // the same row is direct evidence that multiple independent passes
        // actually landed on it, not just a coincidence of one pass hitting
        // a multi-column noise_type like "typo".
        let dir = scratch_dir("hell_multipass");
        let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
        let mut ctx = Context::new(
            "kyc",
            "en",
            &manifest_dir().join("assets/pools").to_string_lossy(),
        )
        .expect("load context");
        let config = build_pipeline_config(
            "kyc",
            5000,
            7,
            "hell",
            0.1,
            0.3,
            &schema,
            "pipeline_test_hell_multipass",
            "ipc",
            false,
            "ipc",
            None,
        )
        .expect("build config");
        ctx.enable_watermark(&config.domain, config.size, config.seed);
        let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");

        let read_all = |path: &str| -> RecordBatch {
            let file = std::fs::File::open(path).unwrap();
            let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
            let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
            let schema = batches[0].schema();
            let refs: Vec<Vec<&dyn Array>> = (0..schema.fields().len())
                .map(|i| batches.iter().map(|b| b.column(i).as_ref()).collect())
                .collect();
            let cols: Vec<ArrayRef> = refs
                .iter()
                .map(|r| arrow::compute::concat(r).unwrap())
                .collect();
            RecordBatch::try_new(schema, cols).unwrap()
        };
        let ds = read_all(&output.output_files[0]);
        let gt = read_all(&output.gt_file);

        let ds_rid = ds.column_by_name("record_id").unwrap().as_string::<i32>();
        let ds_mid = ds.column_by_name("master_id").unwrap().as_string::<i32>();
        let gt_rid = gt.column_by_name("record_id").unwrap().as_string::<i32>();
        // `match_type` is `Dictionary(Int32, Utf8)` — cast
        // back to plain `Utf8` for this test-only comparison.
        let gt_mt_col =
            arrow::compute::cast(gt.column_by_name("match_type").unwrap(), &DataType::Utf8)
                .unwrap();
        let gt_mt = gt_mt_col.as_string::<i32>();

        let mut match_type_by_rid: HashMap<&str, &str> = HashMap::new();
        for i in 0..gt.num_rows() {
            match_type_by_rid.insert(gt_rid.value(i), gt_mt.value(i));
        }

        // The true master row for a given master_id is whichever row was
        // written *first* (base rows are always written before any of their
        // duplicate copies, within a batch) — not "whichever row happens to
        // be labeled `exact_dup`": a cluster with only fuzzy copies has no
        // `exact_dup` row at all (a base row only reads `exact_dup` when the
        // cluster actually contains a byte-identical copy — see `gt.rs`).
        let mut master_row_by_mid: HashMap<&str, usize> = HashMap::new();
        for i in 0..ds.num_rows() {
            master_row_by_mid.entry(ds_mid.value(i)).or_insert(i);
        }

        let compare_cols: Vec<usize> = (0..ds.num_columns())
            .filter(|&i| {
                let field = ds.schema().field(i).clone();
                matches!(field.data_type(), DataType::Utf8)
                    && !matches!(
                        field.name().as_str(),
                        "record_id" | "domain" | "entity_type" | "master_id"
                    )
                    && !field.name().ends_with("_id")
            })
            .collect();

        let mut max_changed_cols = 0usize;
        for i in 0..ds.num_rows() {
            if match_type_by_rid.get(ds_rid.value(i)) != Some(&"fuzzy_dup") {
                continue;
            }
            let Some(&master_i) = master_row_by_mid.get(ds_mid.value(i)) else {
                continue;
            };
            let mut changed = 0usize;
            for &col_i in &compare_cols {
                let col = ds.column(col_i).as_string::<i32>();
                let a = (!col.is_null(i)).then(|| col.value(i));
                let b = (!col.is_null(master_i)).then(|| col.value(master_i));
                if a != b {
                    changed += 1;
                }
            }
            max_changed_cols = max_changed_cols.max(changed);
        }

        assert!(
            max_changed_cols >= 2,
            "expected at least one fuzzy_dup row to differ from its master on \
             2+ columns (evidence of multiple noise passes), got max {max_changed_cols}"
        );

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_record_id_and_pad_string_roundtrip() {
        let rid = record_id_string(7);
        assert_eq!(rid, "R-0000000000007");
        let pad = pad_string(7);
        assert_eq!(pad, "0000000000007");
        let strs = record_id_strs(0..3);
        assert_eq!(strs.len(), 3);
        assert_eq!(strs[0], record_id_string(0));
        assert_eq!(strs[2], record_id_string(2));
    }

    #[test]
    fn test_parse_record_idx_roundtrip() {
        for i in [0usize, 7, 1_234_567, 300_000_000] {
            assert_eq!(parse_record_idx(&record_id_string(i)), Some(i as u64));
        }
        assert_eq!(parse_record_idx("garbage"), None);
        assert_eq!(parse_record_idx("R-1"), None);
        assert_eq!(parse_record_idx("X-0000000000007"), None);
    }

    #[test]
    fn test_pack_master_key_roundtrip_and_ordering() {
        let mid = |e: usize, i: usize| format!("{}-{}", entity_prefix(e), pad_string(i));
        assert!(pack_master_key(&mid(0, 0)).is_some());
        // Same entity, ascending local index -> ascending packed key.
        assert!(pack_master_key(&mid(0, 1)).unwrap() > pack_master_key(&mid(0, 0)).unwrap());
        // Different entity always outranks any local index within a lower
        // entity, exactly like comparing the original fixed-width strings.
        assert!(
            pack_master_key(&mid(1, 0)).unwrap()
                > pack_master_key(&mid(0, 9_999_999_999_999)).unwrap()
        );
        assert_eq!(pack_master_key("garbage"), None);
        assert_eq!(pack_master_key("HN-0000000000007"), None);
        assert_eq!(pack_master_key("CANARY-000-0-abc"), None);
    }

    // ── Isolated measurement: zip cost vs column count ─────────
    //
    // Not a correctness test -- #[ignore]'d so it never runs in normal
    // `cargo test`. Measures whether `apply_noise_with_retry`'s retry-zip
    // loop and `apply_extra_pass_per_row`'s pass-zip loop (both iterate
    // `0..num_columns()`, not just the noise_type's actual target_cols) cost
    // scales with TOTAL column count C, holding the number of noised
    // columns T=1 constant. Run with:
    //   cargo test --release pipeline::tests::bench_zip_cost_vs_column_count -- --ignored --nocapture
    #[test]
    #[ignore]
    fn bench_zip_cost_vs_column_count() {
        use arrow::array::StringArray;
        use arrow::datatypes::{DataType, Field, Schema};

        fn make_batch(n_rows: usize, n_cols: usize) -> RecordBatch {
            let fields: Vec<Field> = (0..n_cols)
                .map(|i| Field::new(format!("col{i}"), DataType::Utf8, false))
                .collect();
            let cols: Vec<ArrayRef> = (0..n_cols)
                .map(|_| {
                    Arc::new(StringArray::from(
                        (0..n_rows)
                            .map(|_| "abcdefghijklmnop".to_string())
                            .collect::<Vec<_>>(),
                    )) as ArrayRef
                })
                .collect();
            RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
        }

        const N_ROWS: usize = 100_000;
        const N_REPS: usize = 5;
        // 3 noise_passes total (hell tier): 1 first pass + 2 extra passes,
        // matching `apply_noise_with_retry` + `apply_extra_pass_per_row`
        // call pattern in `run_pipeline`'s dup-generation closure.
        const EXTRA_PASSES: usize = 2;

        eprintln!("\n[bench_zip] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
        eprintln!("[bench_zip] C=total columns, T=1 noised column (fixed) -- hypothesis: cost ∝ C");

        for &n_cols in &[8usize, 15, 23] {
            let plan_cols = vec!["col0".to_string()];
            let all_types: Vec<&str> = vec!["typo"];
            let mut total = std::time::Duration::ZERO;

            for rep in 0..N_REPS {
                let batch = make_batch(N_ROWS, n_cols);
                let mut rng = Rng::new(42 + rep as u64);
                let t0 = std::time::Instant::now();

                let (noisy, mut unchanged, _) =
                    apply_noise_with_retry(&batch, "typo", &plan_cols, &mut rng, &[]).unwrap();
                let mut current = noisy;
                for _ in 0..EXTRA_PASSES {
                    current = apply_extra_pass_per_row(
                        &current,
                        &mut unchanged,
                        &all_types,
                        &mut rng,
                        &[],
                    )
                    .unwrap();
                }
                std::hint::black_box(&current);
                total += t0.elapsed();
            }

            let avg = total / N_REPS as u32;
            eprintln!(
                "[bench_zip] C={n_cols:2}  avg={:>8.3}ms  per_column={:>7.4}ms  per_row_per_column={:>10.6}us",
                avg.as_secs_f64() * 1000.0,
                avg.as_secs_f64() * 1000.0 / n_cols as f64,
                avg.as_secs_f64() * 1_000_000.0 / (N_ROWS * n_cols) as f64,
            );
        }
    }

    // ── Isolated measurement: kyc residual factor ───────────────────────
    //
    // The `bench_zip_cost_*` measurement above got two things wrong when
    // trying to explain kyc/hell's real-world ~4-5x slowdown vs nonprofit:
    // (1) it assumed `noise_passes = 3` at `hell` -- the real value is
    // `DifficultySettings::passes = 8` (schema.rs); (2) it used generic
    // `col0..colN` column names with `plan_cols=[]` for the "extra pass"
    // calls, but `noise_type_targets_column` is name-based -- generic names
    // match none of hell's 9 categories, so those extra passes were
    // silently no-ops regardless of `EXTRA_PASSES`. This measurement uses
    // the REAL kyc.natural_person (23 cols) and nonprofit.donation (7 cols,
    // the weight-dominant nonprofit entity, weight=200) column NAMES loaded
    // from `schemas/*.json`, hell's real 9-category `noise_types` list, and
    // the real pass count (8) -- reproducing genuine per-category column
    // matching via `noise_type_targets_column`. Run with:
    //   cargo test --release pipeline::tests::bench_residual_kyc_vs_nonprofit -- --ignored --nocapture
    #[test]
    #[ignore]
    fn bench_residual_kyc_vs_nonprofit() {
        use arrow::array::StringArray;
        use arrow::datatypes::{DataType, Field, Schema};

        fn make_batch(n_rows: usize, col_names: &[&str]) -> RecordBatch {
            let fields: Vec<Field> = col_names
                .iter()
                .map(|n| Field::new(*n, DataType::Utf8, true))
                .collect();
            let cols: Vec<ArrayRef> = col_names
                .iter()
                .map(|_| {
                    Arc::new(StringArray::from(
                        (0..n_rows)
                            .map(|_| "abcdefghijklmnop testword".to_string())
                            .collect::<Vec<_>>(),
                    )) as ArrayRef
                })
                .collect();
            RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
        }

        const N_ROWS: usize = 100_000;
        const N_REPS: usize = 5;
        const HELL_TYPES: &[&str] = &[
            "typo_aggressive",
            "visual",
            "unicode_pollution",
            "names",
            "dates",
            "identifiers",
            "addresses",
            "companies",
            "extra",
        ];
        const HELL_PASSES: usize = 8; // schema.rs DifficultySettings::passes for "hell"

        let kyc_cols: &[&str] = &[
            "party_id",
            "given_name",
            "family_name",
            "middle_name",
            "birth_date",
            "birth_place",
            "nationality",
            "email_address",
            "mobile_phone",
            "personal_administrative_number",
            "tax_id",
            "residential_address_line1",
            "residential_city",
            "residential_postal_code",
            "residential_state",
            "residential_country",
            "occupation",
            "risk_score",
            "document_type",
            "document_number",
            "source_system",
        ];
        let nonprofit_cols: &[&str] = &[
            "donation_id",
            "donor_id",
            "campaign_id",
            "donation_date",
            "payment_method",
            "source_system",
        ];

        eprintln!("\n[bench_residual] N_ROWS={N_ROWS} N_REPS={N_REPS} HELL_PASSES={HELL_PASSES}");

        // Per-category breakdown for kyc only (the slow one) -- one rep,
        // timing each of the 9 categories' apply_noise_with_retry call
        // individually (fresh copy of the batch each time, not chained) to
        // find which category(ies) dominate the 18s total.
        {
            let batch = make_batch(N_ROWS, kyc_cols);
            eprintln!("[bench_residual] per-category breakdown on kyc.natural_person:");
            for &ttype in HELL_TYPES {
                let mut rng = Rng::new(42);
                let t0 = std::time::Instant::now();
                let (_, _, tcols) =
                    apply_noise_with_retry(&batch, ttype, &[], &mut rng, &[]).unwrap();
                eprintln!(
                    "[bench_residual]   {ttype:20} {:>8.3}ms  target_cols={:?}",
                    t0.elapsed().as_secs_f64() * 1000.0,
                    tcols,
                );
            }
        }

        for (label, col_names) in [
            ("kyc.natural_person", kyc_cols),
            ("nonprofit.donation", nonprofit_cols),
        ] {
            let mut total = std::time::Duration::ZERO;
            let mut total_target_hits = 0usize;

            for rep in 0..N_REPS {
                let batch = make_batch(N_ROWS, col_names);
                let mut rng = Rng::new(42 + rep as u64);
                let t0 = std::time::Instant::now();

                // First pass: same as real generation, one explicit
                // noise_type (heuristic-matched via plan_cols=[]).
                let (noisy, mut unchanged, t0cols) =
                    apply_noise_with_retry(&batch, "identifiers", &[], &mut rng, &[]).unwrap();
                total_target_hits += t0cols.len();
                let mut current = noisy;
                for _ in 1..HELL_PASSES {
                    current = apply_extra_pass_per_row(
                        &current,
                        &mut unchanged,
                        HELL_TYPES,
                        &mut rng,
                        &[],
                    )
                    .unwrap();
                }
                std::hint::black_box(&current);
                total += t0.elapsed();
            }

            let avg = total / N_REPS as u32;
            eprintln!(
                "[bench_residual] {label:20} C={:2}  avg={:>8.3}ms  avg_first_pass_target_cols={:.1}",
                col_names.len(),
                avg.as_secs_f64() * 1000.0,
                total_target_hits as f64 / N_REPS as f64,
            );
        }
    }

    #[test]
    fn test_chunk_sizes_fold_small_remainder() {
        assert_eq!(chunk_sizes(1000, 500), vec![500, 500]);
        assert_eq!(chunk_sizes(1005, 500), vec![500, 505]);
        assert_eq!(chunk_sizes(1010, 500), vec![500, 500, 10]);
        assert_eq!(chunk_sizes(1234, 400), vec![400, 400, 400, 34]);
        for (total, cs) in [(1005, 500), (99, 10), (10_009, 1000), (25, 12)] {
            let sizes = chunk_sizes(total, cs);
            assert_eq!(sizes.iter().sum::<usize>(), total);
            assert!(sizes.iter().all(|&s| s >= MIN_RUN_SIZE), "{sizes:?}");
        }
    }
}