dupehell_core/
lib.rs

1// DupeHell -- MIT License . Educational Use Only
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// EDUCATIONAL AND RESEARCH PURPOSES ONLY -- see ETHICS.md for prohibited uses.
5// No liability for misuse.
6
7#[cfg(feature = "python")]
8use pyo3::exceptions::PyValueError;
9#[cfg(feature = "python")]
10use pyo3::prelude::*;
11
12mod buf_gen;
13mod canary;
14mod column_gen;
15pub mod context;
16pub mod difficulty;
17mod entity_gen;
18mod fast_template;
19mod fk_remap;
20mod graph_gen;
21pub mod gt;
22mod hn_common;
23mod noise;
24pub mod pipeline;
25mod pool_lookup;
26pub mod rng;
27pub mod schema;
28
29pub use context::Context;
30pub use difficulty::DifficultyReport;
31pub use pipeline::{PipelineConfig, PipelineOutput, PipelineStats, run_pipeline};
32pub use schema::{
33    DomainSchema, EntitySchema, HnSchema, build_pipeline_config, chrono_now, load_schema,
34};
35
36#[cfg(feature = "python")]
37#[pyclass]
38#[derive(Clone)]
39pub struct GenerateResult {
40    #[pyo3(get)]
41    pub dataset: String,
42    #[pyo3(get)]
43    pub ground_truth: String,
44    #[pyo3(get)]
45    pub total_records: usize,
46    #[pyo3(get)]
47    pub exact_dups: usize,
48    #[pyo3(get)]
49    pub hard_negs: usize,
50    #[pyo3(get)]
51    pub uniques: usize,
52    #[pyo3(get)]
53    pub masters: usize,
54    #[pyo3(get)]
55    pub nodes: Option<String>,
56    #[pyo3(get)]
57    pub edges: Option<String>,
58}
59
60#[cfg(feature = "python")]
61#[pymethods]
62impl GenerateResult {
63    fn __repr__(&self) -> String {
64        format!(
65            "GenerateResult(dataset={:?}, ground_truth={:?}, total_records={}, exact_dups={}, hard_negs={}, uniques={}, masters={})",
66            self.dataset,
67            self.ground_truth,
68            self.total_records,
69            self.exact_dups,
70            self.hard_negs,
71            self.uniques,
72            self.masters,
73        )
74    }
75
76    fn __str__(&self) -> String {
77        self.__repr__()
78    }
79}
80
81#[cfg(feature = "python")]
82#[pyfunction]
83fn estimate_difficulty(
84    domain: &str,
85    size: usize,
86    seed: u64,
87    difficulty: &str,
88    schemas_dir: &str,
89    hard_neg_ratio: f64,
90) -> PyResult<String> {
91    let schema =
92        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
93    let report = crate::difficulty::estimate_difficulty(
94        domain,
95        size,
96        seed,
97        difficulty,
98        hard_neg_ratio,
99        &schema,
100    )
101    .map_err(PyValueError::new_err)?;
102    serde_json::to_string(&report).map_err(|e| PyValueError::new_err(e.to_string()))
103}
104
105#[cfg(feature = "python")]
106#[pyfunction]
107#[allow(clippy::too_many_arguments)]
108fn generate(
109    domain: &str,
110    size: usize,
111    seed: u64,
112    difficulty: &str,
113    output_dir: &str,
114    locale: &str,
115    pools_dir: &str,
116    schemas_dir: &str,
117    output_format: &str,
118    hard_neg_ratio: f64,
119    singleton_master_fraction: f64,
120    generate_graph: bool,
121    graph_format: &str,
122) -> PyResult<GenerateResult> {
123    let schema =
124        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
125
126    if output_format != "ipc" && output_format != "parquet" {
127        return Err(PyValueError::new_err(format!(
128            "invalid output format '{output_format}'; expected 'ipc' or 'parquet'"
129        )));
130    }
131
132    let mut ctx = Context::new(domain, locale, pools_dir).map_err(PyValueError::new_err)?;
133
134    let run_id = schema::deterministic_run_id(domain, size, seed, difficulty, hard_neg_ratio);
135    let config = build_pipeline_config(
136        domain,
137        size,
138        seed,
139        difficulty,
140        hard_neg_ratio,
141        singleton_master_fraction,
142        &schema,
143        &run_id,
144        output_format,
145        generate_graph,
146        graph_format,
147    )
148    .map_err(PyValueError::new_err)?;
149
150    ctx.enable_watermark(&config.domain, config.size, config.seed);
151
152    let output = run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?;
153
154    Ok(GenerateResult {
155        dataset: output.output_files.into_iter().next().unwrap_or_default(),
156        ground_truth: output.gt_file,
157        total_records: output.stats.total_records,
158        exact_dups: output.stats.exact_dups,
159        hard_negs: output.stats.hard_negs,
160        uniques: output.stats.uniques,
161        masters: output.stats.masters,
162        nodes: output.nodes,
163        edges: output.edges,
164    })
165}
166
167#[cfg(feature = "python")]
168#[pymodule]
169fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
170    m.add_function(wrap_pyfunction!(generate, m)?)?;
171    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
172    m.add_class::<GenerateResult>()?;
173    Ok(())
174}
dupehell_core/lib.rs

dupehell_core/
lib.rs
