dupehell_core/
lib.rs

1// DupeHell -- MIT License
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 cpu_affinity;
17pub mod difficulty;
18mod entity_gen;
19mod fast_template;
20mod fk_remap;
21mod graph_gen;
22pub mod gt;
23mod hn_common;
24mod noise;
25pub mod pipeline;
26mod pool_lookup;
27pub mod rng;
28pub mod schema;
29
30pub use context::Context;
31pub use difficulty::DifficultyReport;
32pub use pipeline::{
33    PipelineConfig, PipelineOutput, PipelineStats, run_pipeline, run_pipeline_with_progress,
34};
35pub use schema::{
36    DomainSchema, EntitySchema, HnSchema, build_pipeline_config, chrono_now, load_schema,
37};
38
39// `GenerateResult` is only ever constructed in Rust and returned to Python
40// (never passed back in as an argument), so it doesn't need the
41// `FromPyObject` impl pyo3 0.29 now makes opt-in for `Clone` pyclasses.
42#[cfg(feature = "python")]
43#[pyclass(skip_from_py_object)]
44#[derive(Clone)]
45pub struct GenerateResult {
46    #[pyo3(get)]
47    pub dataset: String,
48    #[pyo3(get)]
49    pub ground_truth: String,
50    #[pyo3(get)]
51    pub total_records: usize,
52    #[pyo3(get)]
53    pub exact_dups: usize,
54    #[pyo3(get)]
55    pub fuzzy_dups: usize,
56    #[pyo3(get)]
57    pub hard_negs: usize,
58    #[pyo3(get)]
59    pub uniques: usize,
60    #[pyo3(get)]
61    pub masters: usize,
62    #[pyo3(get)]
63    pub nodes: Option<String>,
64    #[pyo3(get)]
65    pub edges: Option<String>,
66}
67
68#[cfg(feature = "python")]
69#[pymethods]
70impl GenerateResult {
71    fn __repr__(&self) -> String {
72        format!(
73            "GenerateResult(dataset={:?}, ground_truth={:?}, total_records={}, exact_dups={}, fuzzy_dups={}, hard_negs={}, uniques={}, masters={})",
74            self.dataset,
75            self.ground_truth,
76            self.total_records,
77            self.exact_dups,
78            self.fuzzy_dups,
79            self.hard_negs,
80            self.uniques,
81            self.masters,
82        )
83    }
84
85    fn __str__(&self) -> String {
86        self.__repr__()
87    }
88}
89
90/// The chosen difficulty tier's own singleton-master fraction (0.50/0.30/
91/// 0.20/0.10 for light/medium/hard/hell). Exposed so the Python `generate()`
92/// wrapper can derive its default from `difficulty` instead of hardcoding a
93/// single fixed value that silently ignores the tier — see
94/// `schema::default_singleton_master_fraction` for the single source of
95/// truth this mirrors.
96#[cfg(feature = "python")]
97#[pyfunction]
98fn default_singleton_master_fraction(difficulty: &str) -> f64 {
99    schema::default_singleton_master_fraction(difficulty)
100}
101
102#[cfg(feature = "python")]
103#[pyfunction]
104fn estimate_difficulty(
105    domain: &str,
106    size: usize,
107    seed: u64,
108    difficulty: &str,
109    schemas_dir: &str,
110    hard_neg_ratio: f64,
111) -> PyResult<String> {
112    let schema =
113        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
114    let report = crate::difficulty::estimate_difficulty(
115        domain,
116        size,
117        seed,
118        difficulty,
119        hard_neg_ratio,
120        &schema,
121    )
122    .map_err(PyValueError::new_err)?;
123    serde_json::to_string(&report).map_err(|e| PyValueError::new_err(e.to_string()))
124}
125
126#[cfg(feature = "python")]
127#[pyfunction]
128#[pyo3(signature = (
129    domain, size, seed, difficulty, output_dir, locale, pools_dir, schemas_dir,
130    output_format, hard_neg_ratio, singleton_master_fraction, generate_graph,
131    graph_format, only_entity=None,
132))]
133#[allow(clippy::too_many_arguments)]
134fn generate(
135    domain: &str,
136    size: usize,
137    seed: u64,
138    difficulty: &str,
139    output_dir: &str,
140    locale: &str,
141    pools_dir: &str,
142    schemas_dir: &str,
143    output_format: &str,
144    hard_neg_ratio: f64,
145    singleton_master_fraction: f64,
146    generate_graph: bool,
147    graph_format: &str,
148    only_entity: Option<&str>,
149) -> PyResult<GenerateResult> {
150    let schema =
151        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
152
153    if output_format != "ipc" && output_format != "parquet" {
154        return Err(PyValueError::new_err(format!(
155            "invalid output format '{output_format}'; expected 'ipc' or 'parquet'"
156        )));
157    }
158
159    let mut ctx = Context::new(domain, locale, pools_dir).map_err(PyValueError::new_err)?;
160
161    let run_id = schema::deterministic_run_id(
162        domain,
163        size,
164        seed,
165        difficulty,
166        hard_neg_ratio,
167        singleton_master_fraction,
168        locale,
169        only_entity,
170    );
171    let config = build_pipeline_config(
172        domain,
173        size,
174        seed,
175        difficulty,
176        hard_neg_ratio,
177        singleton_master_fraction,
178        &schema,
179        &run_id,
180        output_format,
181        generate_graph,
182        graph_format,
183        only_entity,
184    )
185    .map_err(PyValueError::new_err)?;
186
187    ctx.enable_watermark(&config.domain, config.size, config.seed);
188
189    let output = run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?;
190
191    Ok(GenerateResult {
192        dataset: output.output_files.into_iter().next().unwrap_or_default(),
193        ground_truth: output.gt_file,
194        total_records: output.stats.total_records,
195        exact_dups: output.stats.exact_dups,
196        fuzzy_dups: output.stats.fuzzy_dups,
197        hard_negs: output.stats.hard_negs,
198        uniques: output.stats.uniques,
199        masters: output.stats.masters,
200        nodes: output.nodes,
201        edges: output.edges,
202    })
203}
204
205#[cfg(feature = "python")]
206#[pymodule]
207fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
208    m.add_function(wrap_pyfunction!(generate, m)?)?;
209    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
210    m.add_function(wrap_pyfunction!(default_singleton_master_fraction, m)?)?;
211    m.add_class::<GenerateResult>()?;
212    Ok(())
213}
dupehell_core/lib.rs

dupehell_core/
lib.rs
