dupehell 0.5.5

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

use std::collections::HashMap;
use std::fs::File;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread::JoinHandle;

use arrow::array::{ArrayRef, Float64Array, Float64Builder, Int32Builder, StringBuilder};

use crate::pipeline::DictValues;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::ipc::writer::FileWriter;
use arrow::record_batch::RecordBatch;

/// Output format for the generated property-graph files.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GraphFormat {
    Ipc,
    Parquet,
}

impl GraphFormat {
    pub fn from_str(s: &str) -> GraphFormat {
        match s {
            "parquet" => GraphFormat::Parquet,
            _ => GraphFormat::Ipc,
        }
    }
}

/// Writes `_nodes.{ext}` directly to the final file.
///
/// The node schema is the pipeline `full_arc` with column 0 (`record_id`)
/// renamed `node_id`; all other columns are kept positionally identical.
///
/// The actual IPC encode+write happens on a dedicated background thread
/// (hunt3108_graph/H1): every batch written to the main dataset file used to
/// also be written here, synchronously, right before the dataset write —
/// doubling the IPC-encoding cost of every batch (measured ~+95-97%
/// overhead in isolation, consistent across schema widths). `write_batch`
/// only rebuilds the batch under the node schema (cheap: `Arc`-cloned
/// columns, no data copy) and hands it to the writer thread over a channel,
/// letting that thread's encode+write overlap with whatever the caller does
/// next (typically: write the same batch to the main dataset file).
///
/// Errors from the background thread surface at `finish()`, not at the
/// `write_batch` call where the failing write actually happened — a
/// deliberate trade-off of this design (the alternative, checking after
/// every send, would need a round-trip per batch and defeat the point of
/// overlapping the I/O). On failure, a run keeps generating and writing to
/// the main dataset until `finish()` is reached, which is more wasted work
/// than the previous synchronous version, but write failures here are
/// expected to be rare/environmental (disk full, permissions) and fatal to
/// the run either way.
pub struct NodeWriter {
    tx: Option<mpsc::SyncSender<RecordBatch>>,
    handle: Option<JoinHandle<Result<(), String>>>,
    schema: Arc<Schema>,
}

/// Bounds how many node batches can be in flight (queued in the channel +
/// one being encoded) before `write_batch` blocks and applies backpressure.
/// An unbounded `mpsc::channel` here would let the queue grow without limit
/// whenever the writer thread falls behind the rest of the pipeline (real at
/// high record counts: measured +80% peak RSS at 100M with `--graph`,
/// vs. +5-7% at 1M-10M, hunt3108_graph/H1 follow-up) -- a small bound caps
/// the extra memory to a handful of batches while still letting the node
/// encode of the previous batch overlap with the caller writing the current
/// one to the main dataset (the actual gain this thread exists for).
const NODE_WRITER_CHANNEL_CAP: usize = 2;

impl NodeWriter {
    /// `path` is the final file (written directly, no draft/rename).
    /// `full_schema` is the pipeline `full_arc` (record_id in column 0).
    /// `metadata` is the `dupehell.*` map copied from the dataset.
    pub fn new(
        path: &str,
        full_schema: &Schema,
        metadata: &HashMap<String, String>,
    ) -> Result<Self, String> {
        let fields: Vec<Field> = full_schema
            .fields()
            .iter()
            .enumerate()
            .map(|(i, f)| {
                if i == 0 {
                    Field::new("node_id", f.data_type().clone(), f.is_nullable())
                        .with_metadata(f.metadata().clone())
                } else {
                    f.as_ref().clone()
                }
            })
            .collect();
        let schema = Arc::new(Schema::new(fields).with_metadata(metadata.clone()));

        let file = File::create(path).map_err(|e| format!("create node file {path}: {e}"))?;
        let mut writer = FileWriter::try_new(file, &schema)
            .map_err(|e| format!("node FileWriter {path}: {e}"))?;

        let (tx, rx) = mpsc::sync_channel::<RecordBatch>(NODE_WRITER_CHANNEL_CAP);
        let handle = std::thread::Builder::new()
            .name("dupehell-node-writer".to_string())
            .spawn(move || -> Result<(), String> {
                for rb in rx {
                    writer
                        .write(&rb)
                        .map_err(|e| format!("write node batch: {e}"))?;
                }
                writer
                    .finish()
                    .map_err(|e| format!("finish node writer: {e}"))
            })
            .map_err(|e| format!("spawn node writer thread: {e}"))?;

        Ok(NodeWriter {
            tx: Some(tx),
            handle: Some(handle),
            schema,
        })
    }

    /// `batch` is a base/dup/hn/canary record batch in `full_arc` layout
    /// (record_id in column 0). Rebuilt positionally with the node schema
    /// (column 0 renamed `node_id`) here, on the caller's thread — cheap,
    /// no data copy — then handed off; the actual IPC write happens on the
    /// background thread, overlapped with whatever the caller does next.
    pub fn write_batch(&mut self, batch: &RecordBatch) -> Result<(), String> {
        let rb = RecordBatch::try_new(self.schema.clone(), batch.columns().to_vec())
            .map_err(|e| format!("rebuild node batch: {e}"))?;
        self.tx
            .as_ref()
            .expect("write_batch called after finish()")
            .send(rb)
            .map_err(|_| {
                "node writer thread terminated early (see finish() for the underlying error)"
                    .to_string()
            })
    }

    /// Closes the channel (signals the background thread to finish writing
    /// and close the file), then joins it and surfaces any write error.
    pub fn finish(mut self) -> Result<(), String> {
        drop(self.tx.take());
        match self.handle.take().expect("finish() called twice").join() {
            Ok(res) => res,
            Err(_) => Err("node writer thread panicked".to_string()),
        }
    }
}

/// `Dictionary(Int32, Utf8)` for `edge_type`/`subtype` — hunt1808/H11, an
/// explicit output-contract change (not bit-identical: the declared column
/// type changes from `Utf8`), applied on explicit user request.
fn low_cardinality_dict_type() -> DataType {
    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
}

fn edge_schema(metadata: &HashMap<String, String>) -> Arc<Schema> {
    Arc::new(
        Schema::new(vec![
            Field::new("source_node_id", DataType::Utf8, false),
            Field::new("target_node_id", DataType::Utf8, false),
            Field::new("edge_type", low_cardinality_dict_type(), false),
            Field::new("subtype", low_cardinality_dict_type(), false),
            Field::new("weight", DataType::Float64, false),
        ])
        .with_metadata(metadata.clone()),
    )
}

/// Writes `_edges.{ext}` directly; flushes in bounded batches.
pub struct EdgeWriter {
    writer: FileWriter<File>,
    schema: Arc<Schema>,
    src_buf: StringBuilder,
    tgt_buf: StringBuilder,
    etype_dict: DictValues,
    etype_keys: Int32Builder,
    subtype_dict: DictValues,
    subtype_keys: Int32Builder,
    weight_buf: Float64Builder,
    count: usize,
}

const EDGE_FLUSH: usize = 100_000;

impl EdgeWriter {
    /// `subtype_dict` must already contain every subtype value this run
    /// will ever push for this writer — the full known set (FK
    /// `source_col`s + HN pattern kinds + `"complete"`/`"spanning_tree"`,
    /// see the call site in `pipeline.rs`) — built once, before any edge is
    /// pushed, and shared across every flushed batch (hunt1808/H11; see
    /// `DictValues`'s doc comment for why IPC output requires this).
    /// `edge_type`'s value set is fixed and universal (`"fk"`/`"hard_neg"`/
    /// `"exact_dup"`/`"fuzzy_dup"`, see `pair_edge_key` in
    /// `push_dup_clusters` and the `push_batch` call sites in
    /// `pipeline.rs`), so it's built here rather than threaded in from
    /// every caller.
    pub fn new(
        path: &str,
        metadata: &HashMap<String, String>,
        subtype_dict: DictValues,
    ) -> Result<Self, String> {
        let schema = edge_schema(metadata);
        let file = File::create(path).map_err(|e| format!("create edge file {path}: {e}"))?;
        let writer = FileWriter::try_new(file, &schema)
            .map_err(|e| format!("edge FileWriter {path}: {e}"))?;
        Ok(EdgeWriter {
            writer,
            schema,
            src_buf: StringBuilder::new(),
            tgt_buf: StringBuilder::new(),
            etype_dict: DictValues::new(["fk", "hard_neg", "exact_dup", "fuzzy_dup"]),
            etype_keys: Int32Builder::new(),
            subtype_dict,
            subtype_keys: Int32Builder::new(),
            weight_buf: Float64Builder::new(),
            count: 0,
        })
    }

    /// Resolve `etype`/`subtype` to their dictionary keys once, for a
    /// caller that pushes many edges sharing few distinct etype/subtype
    /// values (`push_dup_clusters`) — see `push_keys`.
    pub fn etype_key(&self, etype: &str) -> i32 {
        self.etype_dict.key(etype)
    }

    pub fn subtype_key(&self, subtype: &str) -> i32 {
        self.subtype_dict.key(subtype)
    }

    /// Same as [`push`](Self::push), but takes already-resolved
    /// `etype`/`subtype` dictionary keys instead of re-hashing the strings
    /// through `DictValues::key` on every call (perf-hunt hunt0109/H9(a)):
    /// `push_dup_clusters` calls this per edge, where `etype`/`subtype`
    /// only ever take 2 values each across the whole cluster set — resolved
    /// once by the caller via `etype_key`/`subtype_key`.
    pub fn push_keys(
        &mut self,
        src: &str,
        tgt: &str,
        etype_key: i32,
        subtype_key: i32,
        weight: f64,
    ) -> Result<(), String> {
        self.src_buf.append_value(src);
        self.tgt_buf.append_value(tgt);
        self.etype_keys.append_value(etype_key);
        self.subtype_keys.append_value(subtype_key);
        self.weight_buf.append_value(weight);
        self.count += 1;
        if self.count >= EDGE_FLUSH {
            self.flush()?;
        }
        Ok(())
    }

    /// Vectorized alternative to `push()` for callers that already hold a
    /// whole column's worth of `src`/`tgt` node ids as materialized
    /// `ArrayRef`s (FK-remap and hard-negative edges — `pipeline.rs`'s
    /// per-row loops used to re-extract each value from an existing
    /// `StringArray` via `.value(i)` and re-append it through a
    /// `StringBuilder`, plus one `HashMap` lookup per row for `etype`/
    /// `subtype`, even though `etype`/`subtype`/`weight` are constant for
    /// the whole call — hunt3108_graph/H2, measured ~12x faster than the
    /// scalar loop it replaces). `src`/`tgt` are reused via `Arc` clone —
    /// zero-copy, no re-encoding of already-built strings.
    ///
    /// Flushes any buffered scalar rows first: edges have no ordering
    /// semantics downstream (ground truth reads them as an unordered set of
    /// pairs), so this is purely to keep the file a simple append sequence
    /// rather than interleaving a partially-filled scalar buffer with a
    /// direct batch write.
    pub fn push_batch(
        &mut self,
        src: &ArrayRef,
        tgt: &ArrayRef,
        etype: &str,
        subtype: &str,
        weight: f64,
    ) -> Result<(), String> {
        self.flush()?;
        let n = src.len();
        if n == 0 {
            return Ok(());
        }
        let weight_arr: ArrayRef = Arc::new(Float64Array::from(vec![weight; n]));
        let rb = RecordBatch::try_new(
            self.schema.clone(),
            vec![
                src.clone(),
                tgt.clone(),
                self.etype_dict.const_array(etype, n),
                self.subtype_dict.const_array(subtype, n),
                weight_arr,
            ],
        )
        .map_err(|e| format!("build edge batch: {e}"))?;
        self.writer
            .write(&rb)
            .map_err(|e| format!("write edge batch: {e}"))
    }

    fn flush(&mut self) -> Result<(), String> {
        if self.count == 0 {
            return Ok(());
        }
        let rb = RecordBatch::try_new(
            self.schema.clone(),
            vec![
                Arc::new(self.src_buf.finish()) as ArrayRef,
                Arc::new(self.tgt_buf.finish()) as ArrayRef,
                self.etype_dict.finish_keys(self.etype_keys.finish()),
                self.subtype_dict.finish_keys(self.subtype_keys.finish()),
                Arc::new(self.weight_buf.finish()) as ArrayRef,
            ],
        )
        .map_err(|e| format!("build edge batch: {e}"))?;
        self.writer
            .write(&rb)
            .map_err(|e| format!("write edge batch: {e}"))?;
        // `ArrayBuilder::finish` already takes `&mut self` and resets the
        // builder's internal buffers — no need to replace it with a fresh
        // one, that was just throwing away and reallocating empty builders.
        self.count = 0;
        Ok(())
    }

    pub fn finish(mut self) -> Result<(), String> {
        self.flush()?;
        self.writer
            .finish()
            .map_err(|e| format!("finish edge writer: {e}"))
    }
}

/// Emit duplicate-cluster edges. For a cluster of size `k`, emit the full
/// `k(k-1)/2` complete graph unless it exceeds `max_edges`, in which case a
/// deterministic spanning tree (sorted order) is emitted instead. Each
/// edge's `edge_type` reflects whether both endpoints are genuinely
/// byte-identical (`exact_dup`) or at least one was noised (`fuzzy_dup`) —
/// see `pair_edge_key` below.
pub fn push_dup_clusters(
    ew: &mut EdgeWriter,
    clusters: &crate::gt::ClusterCsr,
    max_edges: usize,
) -> Result<(), String> {
    // Resolved once, not once per edge (perf-hunt hunt0109/H9(a)): `etype`
    // only ever takes 2 values here (`exact_dup`/`fuzzy_dup`: `exact_dup`
    // only when *both* ends are byte-for-byte identical to the cluster's
    // master — transitively identical to each other too; `fuzzy_dup` as
    // soon as either end was genuinely noised, since two fuzzy copies
    // aren't guaranteed to match each other exactly) and `subtype` only 2
    // (`complete`/`spanning_tree`), but the previous code re-hashed one of
    // each through
    // `DictValues::key`'s `HashMap<String, i32>` on every single edge —
    // this cluster-edge path is the highest-volume edge source in the
    // graph output (up to k(k-1)/2 edges per duplicated cluster).
    let k_exact = ew.etype_key("exact_dup");
    let k_fuzzy = ew.etype_key("fuzzy_dup");
    let k_complete = ew.subtype_key("complete");
    let k_spanning = ew.subtype_key("spanning_tree");
    let pair_edge_key = |a: bool, b: bool| if a && b { k_exact } else { k_fuzzy };

    // Reused across every edge endpoint (perf-hunt hunt0109/H9(b)) instead
    // of one `String` allocation per endpoint via `record_id_string` — see
    // `pipeline::append_record_id`.
    let mut src_buf: Vec<u8> = Vec::with_capacity(crate::pipeline::RID_LEN);
    let mut tgt_buf: Vec<u8> = Vec::with_capacity(crate::pipeline::RID_LEN);

    for (records, idents) in clusters.groups() {
        let k = records.len();
        if k < 2 {
            continue;
        }
        // `records` is already ascending within its cluster (guaranteed by
        // `ClusterCsr::build`'s sort), matching the previous
        // `sorted.sort_by(|a, b| a.0.cmp(&b.0))` on the record_id string --
        // no re-sort needed here.
        let n_edges = k * (k - 1) / 2;

        if n_edges > max_edges {
            log::warn!(
                "dup cluster has {n_edges} edges > {max_edges}, using spanning tree fallback"
            );
            for i in 0..k - 1 {
                let etype_key = pair_edge_key(idents.get(i), idents.get(i + 1));
                src_buf.clear();
                crate::pipeline::append_record_id(&mut src_buf, records[i]);
                tgt_buf.clear();
                crate::pipeline::append_record_id(&mut tgt_buf, records[i + 1]);
                let src = std::str::from_utf8(&src_buf).expect("append_record_id: ASCII digits");
                let tgt = std::str::from_utf8(&tgt_buf).expect("append_record_id: ASCII digits");
                ew.push_keys(src, tgt, etype_key, k_spanning, 1.0)?;
            }
        } else {
            for i in 0..k {
                src_buf.clear();
                crate::pipeline::append_record_id(&mut src_buf, records[i]);
                let src = std::str::from_utf8(&src_buf).expect("append_record_id: ASCII digits");
                // `j` indexes both `records` and `idents` in lockstep —
                // an iterator/enumerate rewrite would need the same offset
                // arithmetic clippy is suggesting away, for no clarity gain.
                #[allow(clippy::needless_range_loop)]
                for j in (i + 1)..k {
                    let etype_key = pair_edge_key(idents.get(i), idents.get(j));
                    tgt_buf.clear();
                    crate::pipeline::append_record_id(&mut tgt_buf, records[j]);
                    let tgt =
                        std::str::from_utf8(&tgt_buf).expect("append_record_id: ASCII digits");
                    ew.push_keys(src, tgt, etype_key, k_complete, 1.0)?;
                }
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::AsArray;
    use std::collections::HashSet;

    fn read_edges(path: &str) -> Vec<(String, String, String, String, f64)> {
        let file = File::open(path).unwrap();
        let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
        let mut out = Vec::new();
        for b in reader {
            let b = b.unwrap();
            let src = b.column(0).as_string::<i32>();
            let tgt = b.column(1).as_string::<i32>();
            // `edge_type`/`subtype` are `Dictionary(Int32, Utf8)`
            // (hunt1808/H11) — cast back to plain `Utf8` for this test-only
            // comparison.
            let et_col = arrow::compute::cast(b.column(2), &DataType::Utf8).unwrap();
            let st_col = arrow::compute::cast(b.column(3), &DataType::Utf8).unwrap();
            let et = et_col.as_string::<i32>();
            let st = st_col.as_string::<i32>();
            let w = b.column(4).as_primitive::<arrow::datatypes::Float64Type>();
            for i in 0..b.num_rows() {
                out.push((
                    src.value(i).to_string(),
                    tgt.value(i).to_string(),
                    et.value(i).to_string(),
                    st.value(i).to_string(),
                    w.value(i),
                ));
            }
        }
        out
    }

    fn temp_path(name: &str) -> String {
        let mut p = std::env::temp_dir();
        p.push(format!("dupehell_test_{}_{}", name, std::process::id()));
        p.to_string_lossy().to_string()
    }

    /// Alias for `pipeline::record_id_string`, for brevity in fixtures.
    fn rid(i: usize) -> String {
        crate::pipeline::record_id_string(i)
    }

    /// Builds a single-cluster `ClusterCsr` for a test, from a master's
    /// local index (entity index fixed at 0 -- these tests only cover one
    /// entity type) and its `(record_idx, is_identical)` members.
    fn one_cluster(master_local_idx: usize, members: Vec<(usize, bool)>) -> crate::gt::ClusterCsr {
        let master_key = crate::pipeline::pack_master_key(&format!(
            "{}-{}",
            crate::pipeline::entity_prefix(0),
            crate::pipeline::pad_string(master_local_idx)
        ))
        .unwrap();
        let pairs = members
            .into_iter()
            .map(|(ridx, ident)| {
                (
                    master_key,
                    crate::gt::pack_ridx_identical(ridx as u64, ident),
                )
            })
            .collect();
        crate::gt::ClusterCsr::build(pairs)
    }

    #[test]
    fn push_dup_clusters_complete() {
        let path = temp_path("edges_complete.ipc");
        let _ = std::fs::remove_file(&path);
        let mut ew = EdgeWriter::new(
            &path,
            &HashMap::new(),
            DictValues::new(["complete", "spanning_tree"]),
        )
        .unwrap();
        let clusters = one_cluster(1, vec![(1, true), (2, true), (3, true), (4, true)]);
        push_dup_clusters(&mut ew, &clusters, 10_000).unwrap();
        ew.finish().unwrap();

        let edges = read_edges(&path);
        assert_eq!(edges.len(), 6, "4 records -> 6 complete edges");
        assert!(
            edges
                .iter()
                .all(|e| e.2 == "exact_dup" && e.3 == "complete")
        );
        assert!(edges.iter().all(|e| (e.4 - 1.0).abs() < 1e-9));
        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn push_dup_clusters_fuzzy_edge_type() {
        let path = temp_path("edges_fuzzy.ipc");
        let _ = std::fs::remove_file(&path);
        let mut ew = EdgeWriter::new(
            &path,
            &HashMap::new(),
            DictValues::new(["complete", "spanning_tree"]),
        )
        .unwrap();
        // rid(1) (master) and rid(2) stayed identical; rid(3) was genuinely noised.
        let clusters = one_cluster(1, vec![(1, true), (2, true), (3, false)]);
        push_dup_clusters(&mut ew, &clusters, 10_000).unwrap();
        ew.finish().unwrap();

        let edges = read_edges(&path);
        assert_eq!(edges.len(), 3);
        let by_pair: HashMap<(String, String), String> = edges
            .iter()
            .map(|e| ((e.0.clone(), e.1.clone()), e.2.clone()))
            .collect();
        assert_eq!(by_pair[&(rid(1), rid(2))], "exact_dup");
        assert_eq!(by_pair[&(rid(1), rid(3))], "fuzzy_dup");
        assert_eq!(by_pair[&(rid(2), rid(3))], "fuzzy_dup");
        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn push_dup_clusters_spanning_tree() {
        let path = temp_path("edges_spanning.ipc");
        let _ = std::fs::remove_file(&path);
        let mut ew = EdgeWriter::new(
            &path,
            &HashMap::new(),
            DictValues::new(["complete", "spanning_tree"]),
        )
        .unwrap();
        let members: Vec<(usize, bool)> = (0..200).map(|i| (i, true)).collect();
        let clusters = one_cluster(1, members.clone());
        // 200*199/2 = 19900 edges > 10000 -> spanning tree fallback (199 edges)
        push_dup_clusters(&mut ew, &clusters, 10_000).unwrap();
        ew.finish().unwrap();

        let edges = read_edges(&path);
        assert_eq!(
            edges.len(),
            199,
            "200-record cluster -> 199 spanning-tree edges"
        );
        assert!(
            edges
                .iter()
                .all(|e| e.2 == "exact_dup" && e.3 == "spanning_tree")
        );
        // Adjacent ascending pairs only (members are already 0..200 ascending).
        let got: HashSet<(String, String)> =
            edges.iter().map(|e| (e.0.clone(), e.1.clone())).collect();
        for w in members.windows(2) {
            assert!(got.contains(&(rid(w[0].0), rid(w[1].0))));
        }
        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn push_dup_clusters_skips_singletons() {
        let path = temp_path("edges_singleton.ipc");
        let _ = std::fs::remove_file(&path);
        let mut ew = EdgeWriter::new(
            &path,
            &HashMap::new(),
            DictValues::new(["complete", "spanning_tree"]),
        )
        .unwrap();
        let clusters = one_cluster(1, vec![(1, true)]);
        push_dup_clusters(&mut ew, &clusters, 10_000).unwrap();
        ew.finish().unwrap();
        assert!(read_edges(&path).is_empty());
        let _ = std::fs::remove_file(&path);
    }
}