dupehell_core/
lib.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6#[cfg(feature = "python")]
7use pyo3::exceptions::PyValueError;
8#[cfg(feature = "python")]
9use pyo3::prelude::*;
10
11mod buf_gen;
12mod canary;
13mod column_gen;
14pub mod context;
15pub mod cpu_affinity;
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#[pyo3(signature = (
128    domain, size, seed, difficulty, output_dir, locale, pools_dir, schemas_dir,
129    output_format, hard_neg_ratio, singleton_master_fraction, generate_graph,
130    graph_format, only_entity=None, chunk_size=None,
131))]
132#[allow(clippy::too_many_arguments)]
133fn generate(
134    domain: &str,
135    size: usize,
136    seed: u64,
137    difficulty: &str,
138    output_dir: &str,
139    locale: &str,
140    pools_dir: &str,
141    schemas_dir: &str,
142    output_format: &str,
143    hard_neg_ratio: f64,
144    singleton_master_fraction: f64,
145    generate_graph: bool,
146    graph_format: &str,
147    only_entity: Option<&str>,
148    chunk_size: Option<usize>,
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        chunk_size,
171    );
172
173    let output = match chunk_size {
174        Some(cs) if cs > 0 && cs < size => {
175            ctx.enable_watermark(domain, size, seed);
176            crate::pipeline::run_chunked(
177                &ctx,
178                domain,
179                size,
180                cs,
181                seed,
182                difficulty,
183                hard_neg_ratio,
184                singleton_master_fraction,
185                &schema,
186                &run_id,
187                output_format,
188                generate_graph,
189                graph_format,
190                only_entity,
191                output_dir,
192                None,
193            )
194            .map_err(PyValueError::new_err)?
195        }
196        _ => {
197            let config = build_pipeline_config(
198                domain,
199                size,
200                seed,
201                difficulty,
202                hard_neg_ratio,
203                singleton_master_fraction,
204                &schema,
205                &run_id,
206                output_format,
207                generate_graph,
208                graph_format,
209                only_entity,
210            )
211            .map_err(PyValueError::new_err)?;
212            ctx.enable_watermark(&config.domain, config.size, config.seed);
213            run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?
214        }
215    };
216
217    Ok(GenerateResult {
218        dataset: output.output_files.into_iter().next().unwrap_or_default(),
219        ground_truth: output.gt_file,
220        total_records: output.stats.total_records,
221        exact_dups: output.stats.exact_dups,
222        fuzzy_dups: output.stats.fuzzy_dups,
223        hard_negs: output.stats.hard_negs,
224        uniques: output.stats.uniques,
225        masters: output.stats.masters,
226        nodes: output.nodes,
227        edges: output.edges,
228    })
229}
230
231#[cfg(feature = "python")]
232#[pymodule]
233fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
234    m.add_function(wrap_pyfunction!(generate, m)?)?;
235    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
236    m.add_function(wrap_pyfunction!(default_singleton_master_fraction, m)?)?;
237    m.add_class::<GenerateResult>()?;
238    Ok(())
239}
dupehell_core/lib.rs

dupehell_core/
lib.rs
