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

dupehell_core/
lib.rs
