dupehell_core/

pipeline.rs

// DupeHell -- MIT License . Educational Use Only
//
// Synthetic multi-domain dataset generator for record linkage benchmarking.
// EDUCATIONAL AND RESEARCH PURPOSES ONLY -- see ETHICS.md for prohibited uses.
// No liability for misuse.

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

use arrow::array::{
    Array, ArrayBuilder, ArrayRef, AsArray, StringArray, UInt64Array, UInt64Builder,
};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
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,
}

#[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>,
}

#[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,
}

// ── 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.
const RID_LEN: usize = 15; // "R-" + 13 digits
const PAD_LEN: usize = 13; // 13 digits

#[inline]
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")
}

#[inline]
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")
}

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

/// 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.
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: Vec<String>,
}

// ── 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 — see the 40-domain schema audit in the 2026-07 session).
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 (BUGS.md C13). 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 — see the 40-domain schema audit in the 2026-07
/// session).
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` and the `hunt2407.md`
/// perf-hunt entry this fixes (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 (BUGS.md C20).
            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 (BUGS.md C12). `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" => {
            contains_any(&lower, &["date", "birth", "incorporation", "founding"])
        }
        "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",
            ],
        ),
        "extra"
        | "name_null"
        | "dob_null"
        | "blocking_fail_initial"
        | "blocking_fail_partial"
        | "fuzzy_match"
        | "phonetic" => {
            // Same exclusions as the typo/visual arm above (BUGS.md C12/C20):
            // 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;
            }
            contains_any(&lower, &["address", "street", "city", "note", "comment"])
                || 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,
            // BUGS.md C20). `email_address` is the same collision: it
            // contains "address" but is an email, not a postal address —
            // `identifiers`/`corrupt_email` already handles it (perf-hunt
            // hunt2407.md, 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`.
            !lower.contains("ip_address")
                && !lower.contains("email")
                && contains_any(
                    &lower,
                    &["address", "street", "postal", "city", "state", "country"],
                )
        }
        "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();
    let mut changed = vec![false; n];
    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>();
        for (i, is_changed) in changed.iter_mut().enumerate() {
            if *is_changed {
                continue;
            }
            let av = (!a.is_null(i)).then(|| a.value(i));
            let bv = (!b.is_null(i)).then(|| b.value(i));
            if av != bv {
                *is_changed = true;
            }
        }
    }
    Ok(changed.into_iter().map(|c| !c).collect())
}

/// 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` (perf-hunt H1,
/// hunt2407.md): `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(unchanged.clone());
        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}"))?;
        unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
        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 (perf-hunt H4, hunt2407.md),
    // 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 — hunt2407.md), 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.
    let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
    for (i, &t) in per_row_type.iter().enumerate() {
        groups[t].push(i as u32);
    }

    let schema = noisy.schema();
    let mut pass_unchanged = vec![true; n];
    // column name -> full-length scatter buffer, lazily seeded from
    // `noisy`'s own values and overwritten in place per group that
    // touches it. Only columns some active type actually targets get an
    // entry — untouched columns are left as direct clones of `noisy`'s
    // arrays at the end, same "don't rebuild what didn't change" principle
    // as the H1 fix above.
    let mut touched: HashMap<String, Vec<Option<String>>> = HashMap::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)?;
        for col_name in &target_cols {
            let Ok(col_idx) = schema.index_of(col_name) else {
                continue;
            };
            let noised_col = noised_subset.column(col_idx).as_string::<i32>();
            let buf = touched.entry(col_name.clone()).or_insert_with(|| {
                let orig = noisy.column(col_idx).as_string::<i32>();
                (0..n)
                    .map(|i| (!orig.is_null(i)).then(|| orig.value(i).to_string()))
                    .collect()
            });
            for (local_pos, &orig_i) in idxs.iter().enumerate() {
                buf[orig_i as usize] = (!noised_col.is_null(local_pos))
                    .then(|| noised_col.value(local_pos).to_string());
            }
        }
        for (local_pos, &orig_i) in idxs.iter().enumerate() {
            pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
        }
    }

    let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
    for (col_name, buf) in touched {
        if let Ok(col_idx) = schema.index_of(&col_name) {
            merged_cols[col_idx] = Arc::new(StringArray::from(buf));
        }
    }
    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: &HashMap<String, String>,
    ) -> 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: Vec<parquet::file::metadata::KeyValue> = metadata
                .iter()
                .map(|(k, v)| parquet::file::metadata::KeyValue {
                    key: k.clone(),
                    value: Some(v.clone()),
                })
                .collect();
            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}")),
        }
    }
}

// ── 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,
    mut progress: Option<&mut dyn FnMut(usize, usize)>,
) -> Result<PipelineOutput, String> {
    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();

    // ── 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));
    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: fed one batch at a time as the dataset is generated
    // (see `crate::gt::GtAccumulator`) instead of accumulating full-dataset
    // arrays in RAM.
    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 gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
    let mut gt_acc = crate::gt::GtAccumulator::new(&gt_draft_path)?;

    let mut global_rid_offset = 0usize;

    // ── 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)
                .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.
    const FK_POOL_CAP: usize = 200_000;

    // 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.
    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 col_json_str = &plan.columns_json;

        // Per-entity streaming state
        let fk_targeted = fk_targets.contains(plan.name.as_str());
        let mut fk_builder = arrow::array::StringBuilder::new();
        let mut fk_rid_builder = arrow::array::StringBuilder::new();
        let mut fk_count: usize = 0;
        // 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<Vec<String>> = Vec::new();
        let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
        let mut fk_rng = Rng::new(config.seed.wrapping_add(42));

        // 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();

        // 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: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
        let mut const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();

        // 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 = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);

            // Generate batch
            let request_json = format!(
                r#"{{"entity_name":"{}","n":{},"seed":{},"columns":{},"row_offset":{}}}"#,
                plan.name, batch_n, batch_seed, col_json_str, offset,
            );
            let rb = crate::entity_gen::generate_entity_batch(ctx, &request_json)?;

            // FK remap (uses pools from previously processed entities)
            let mut fk_edges_by_remap: Vec<(String, Vec<String>, 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) {
                        let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
                            &r,
                            pool,
                            &remap.source_col,
                            &mut fk_rng,
                            config.graph_enabled,
                        )?;
                        r = remapped;
                        if config.graph_enabled
                            && let Some(rids) = target_rids
                        {
                            fk_edges_by_remap.push((
                                remap.source_col.clone(),
                                rids,
                                remap.fk_weight,
                            ));
                        }
                    }
                }
                r
            } else {
                rb
            };

            // Extract FK identifiers for this entity's pool, capped at
            // FK_POOL_CAP (see comment above).
            if fk_targeted
                && fk_count < FK_POOL_CAP
                && 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 fk_count >= FK_POOL_CAP {
                        break;
                    }
                    if !s.is_null(i) {
                        fk_builder.append_value(s.value(i));
                        if config.graph_enabled {
                            fk_rid_builder.append_value(record_id_string(global_rid_offset + i));
                        }
                        fk_count += 1;
                    }
                }
            }

            // 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_strs(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_mids: Vec<String> = (offset..offset + batch_n)
                .map(|i| format!("{}-{}", prefix, pad_string(i)))
                .collect();
            let batch_mids_slice = &batch_mids[..];

            // Add metadata + IPC write
            let rid_slice = record_id_strs(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,
                &rid_slice,
                batch_mids_slice,
                &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}"))?;
                for i in 0..batch_n {
                    for (subtype, target_rids, weight) in &fk_edges_by_remap {
                        ew.push(&rid_slice[i], &target_rids[i], "fk", subtype, *weight)
                            .map_err(|e| format!("push fk edge: {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 collect
            gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;

            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_mids_buf: Vec<String> = Vec::new();

                // Collect FK columns to exclude from noise
                let fk_exclude_cols: Vec<String> = plan
                    .fk_remaps
                    .iter()
                    .map(|r| r.source_col.clone())
                    .collect();

                // 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, used by
                // additional noise passes (below) to independently draw a type
                // to apply on top of the first pass's result.
                let all_types: Vec<&str> = plan
                    .noise_types
                    .iter()
                    .map(|n| n.noise_type.as_str())
                    .collect();
                let all_types_ref = &all_types;
                let noise_passes = config.noise_passes.max(1);

                type DupNoiseResult = Result<(RecordBatch, Vec<String>, Vec<bool>), String>;
                let mut results: Vec<DupNoiseResult> = Vec::with_capacity(ndata.len());
                let fk_exclude_ref = &fk_exclude_cols;
                let rb_ref = &rb;
                std::thread::scope(|s| {
                    let mut handles = Vec::with_capacity(ndata.len());
                    for (indices, seed, ntype, cols, cnt) in &ndata {
                        // `thread::scope` lets each spawned closure borrow
                        // `rb`/`indices`/`cols`/`fk_exclude_cols` directly —
                        // the scope guarantees every thread joins before these
                        // borrows end, so the per-thread `.clone()` of the
                        // batch, index array, column list and exclude list
                        // (all cheap Arc/Vec clones, but still real allocations
                        // times the noise-type count) was unnecessary.
                        let prefix_ref: &str = &prefix;
                        handles.push(s.spawn(move || {
                            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 = offset + indices.value(j) as usize;
                                mb.push(format!("{}-{}", prefix_ref, pad_string(global_idx)));
                            }
                            Ok((noisy, mb, unchanged))
                        }));
                    }
                    for h in handles {
                        results.push(h.join().unwrap());
                    }
                });
                let mut dup_is_identical_buf: Vec<bool> = Vec::new();
                for res in results {
                    let (rb, mb, unchanged) = res?;
                    dup_batches.push(rb);
                    dup_mids_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_mids_buf.len();
                    let dup_rids = record_id_strs(global_rid_offset..global_rid_offset + 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,
                        &dup_rids,
                        &dup_mids_buf,
                        &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}"))?;
                    }

                    let is_identical_arr: ArrayRef =
                        Arc::new(arrow::array::BooleanArray::from(dup_is_identical_buf));
                    gt_acc.push_dup_batch(
                        dup_rb_full.column(0),
                        dup_rb_full.column(2),
                        dup_rb_full.column(3),
                        &is_identical_arr,
                    )?;
                    global_rid_offset += dup_total;
                }
            }
        }

        // Save FK pool for cross-entity remapping
        if let Some(id_col) = &plan.identifier_col
            && fk_builder.len() > 0
        {
            let arr = Arc::new(fk_builder.finish()) 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. Same
                // bounds as the identifier column above (FK_POOL_CAP gated).
                let rid_arr = Arc::new(fk_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 = if config.graph_enabled {
                hn_slice_rids.into_iter().flatten().collect::<Vec<String>>()
            } else {
                Vec::new()
            };
            hn_pools.insert(
                plan.name.clone(),
                HnPool {
                    batch,
                    total_count,
                    record_ids,
                },
            );
        }
    }
    let t1e_elapsed = t1e.elapsed().as_secs_f64();
    log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");

    // ── 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: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
    let mut hn_const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();
    // 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 = 0;
    for hn_cfg in config.hard_neg_types.iter() {
        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,
            config.seed.wrapping_add(200),
            config.graph_enabled,
        )?;

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

        let hn_rids = record_id_strs(global_rid_offset..global_rid_offset + n_hn);

        let hn_mids: Vec<String> = (0..n_hn)
            .map(|_i| {
                let id = hn_master_id_counter;
                hn_master_id_counter += 1;
                format!("HN-{:09}", id)
            })
            .collect();

        // 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,
            &hn_rids,
            &hn_mids,
            &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 {
                for i in 0..n_hn {
                    let tgt = &pool_data.record_ids[idx_a[i]];
                    ew.push(&hn_rids[i], tgt, "hard_neg", &pattern, 1.0)
                        .map_err(|e| format!("push hn edge: {e}"))?;
                }
            }
        }

        // Phase 13: collect ArrayRefs instead of per-row StringBuilder
        gt_acc.push_other_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();

    // ── Canary records ────────────────────────────────────────────────────
    {
        let mut canary_null_cache: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
        let mut canary_const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();
        crate::canary::generate_all(
            ctx,
            config,
            &full_arc,
            &mut canary_null_cache,
            &mut canary_const_arr_cache,
            &mut global_rid_offset,
            &fk_pools,
            &mut writer,
            &mut node_writer,
            &mut gt_acc,
        )?;
    }

    // ── 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();

    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,
    } = gt_acc.finish(
        config.difficulty.as_str(),
        &config.output_format,
        &gt_path,
        &metadata,
    )?;
    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();

    // ── 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);
    }

    Ok(PipelineOutput {
        output_files: vec![dataset_path],
        gt_file: gt_path,
        stats,
        nodes: graph_nodes_final,
        edges: graph_edges_final,
    })
}

/// 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}"))?;
    let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
        .metadata()
        .iter()
        .map(|(k, v)| parquet::file::metadata::KeyValue {
            key: k.clone(),
            value: Some(v.clone()),
        })
        .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(())
}

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

fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
    let ts = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs().to_string())
        .unwrap_or_else(|_| "unknown".into());
    HashMap::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).
fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
    let mut field_map: Vec<(String, DataType, bool)> = Vec::new();
    let metadata_fields = ["record_id", "domain", "entity_type", "master_id"];
    for mf in &metadata_fields {
        field_map.push((mf.to_string(), DataType::Utf8, false));
    }
    for plan in &config.entity_plans {
        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;
            }
            let col_type = 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(metadata.clone())
}

/// 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,
    record_ids: &[String],
    master_ids: &[String],
    full_arc: &Arc<Schema>,
    col_lookup: &[Option<usize>],
    null_cache: &mut HashMap<(DataType, usize), ArrayRef>,
    const_arr_cache: &mut HashMap<(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 StringArray on every call, same idea as
    // `null_cache` above for the always-null columns.
    let domain_arr = const_arr_cache
        .entry((domain.to_string(), n))
        .or_insert_with(|| {
            Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
                domain, n,
            ))) as ArrayRef
        })
        .clone();
    let et_arr = const_arr_cache
        .entry((entity_type.to_string(), n))
        .or_insert_with(|| {
            Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
                entity_type,
                n,
            ))) as ArrayRef
        })
        .clone();
    let rid_arr = Arc::new(StringArray::from_iter_values(
        record_ids.iter().map(|s| s.as_str()),
    )) as ArrayRef;
    let mid_arr = Arc::new(StringArray::from_iter_values(
        master_ids.iter().map(|s| s.as_str()),
    )) as ArrayRef;

    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 dt = full_arc.field(i + 4).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};

    /// Regression for the perf-hunt cross-contamination bug (hunt2407.md):
    /// `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",
        )
        .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",
        )
        .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");
            let edge_type = batch
                .column_by_name("edge_type")
                .expect("edge_type col")
                .as_any()
                .downcast_ref::<arrow::array::StringArray>()
                .expect("edge_type as string");
            let subtype = batch
                .column_by_name("subtype")
                .expect("subtype col")
                .as_any()
                .downcast_ref::<arrow::array::StringArray>()
                .expect("subtype as string");
            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);
    }

    #[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",
        )
        .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>();
        let gt_mt = gt.column_by_name("match_type").unwrap().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));
    }

    // ── perf-hunt H1 isolated measurement (hunt2407.md) ─────────────────
    //
    // 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 -- H1's hypothesis. Run with:
    //   cargo test --release pipeline::tests::hunt_h1_zip_cost_vs_column_count -- --ignored --nocapture
    #[test]
    #[ignore]
    fn hunt_h1_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[hunt_h1] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
        eprintln!("[hunt_h1] C=total columns, T=1 noised column (fixed) -- H1 predicts 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!(
                "[hunt_h1] 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,
            );
        }
    }

    // ── perf-hunt kyc residual factor (hunt2407.md, corrected) ──────────
    //
    // The first `hunt_h1_*` 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::hunt_h1_residual_kyc_vs_nonprofit -- --ignored --nocapture
    #[test]
    #[ignore]
    fn hunt_h1_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[hunt_h1_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!("[hunt_h1_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!(
                    "[hunt_h1_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!(
                "[hunt_h1_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,
            );
        }
    }
}