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    if !schema::valid_difficulty_names().contains(&difficulty) {
159        return Err(PyValueError::new_err(format!(
160            "invalid difficulty '{difficulty}'; expected one of {:?}",
161            schema::valid_difficulty_names()
162        )));
163    }
164
165    let mut ctx = Context::new(domain, locale, pools_dir).map_err(PyValueError::new_err)?;
166
167    let run_id = schema::deterministic_run_id(
168        domain,
169        size,
170        seed,
171        difficulty,
172        hard_neg_ratio,
173        singleton_master_fraction,
174        locale,
175        only_entity,
176        chunk_size,
177    );
178
179    let output = match chunk_size {
180        Some(cs) if cs > 0 && cs < size => {
181            ctx.enable_watermark(domain, size, seed);
182            crate::pipeline::run_chunked(
183                &ctx,
184                domain,
185                size,
186                cs,
187                seed,
188                difficulty,
189                hard_neg_ratio,
190                singleton_master_fraction,
191                &schema,
192                &run_id,
193                output_format,
194                generate_graph,
195                graph_format,
196                only_entity,
197                output_dir,
198                None,
199            )
200            .map_err(PyValueError::new_err)?
201        }
202        _ => {
203            let config = build_pipeline_config(
204                domain,
205                size,
206                seed,
207                difficulty,
208                hard_neg_ratio,
209                singleton_master_fraction,
210                &schema,
211                &run_id,
212                output_format,
213                generate_graph,
214                graph_format,
215                only_entity,
216            )
217            .map_err(PyValueError::new_err)?;
218            ctx.enable_watermark(&config.domain, config.size, config.seed);
219            run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?
220        }
221    };
222
223    Ok(GenerateResult {
224        dataset: output.output_files.into_iter().next().unwrap_or_default(),
225        ground_truth: output.gt_file,
226        total_records: output.stats.total_records,
227        exact_dups: output.stats.exact_dups,
228        fuzzy_dups: output.stats.fuzzy_dups,
229        hard_negs: output.stats.hard_negs,
230        uniques: output.stats.uniques,
231        masters: output.stats.masters,
232        nodes: output.nodes,
233        edges: output.edges,
234    })
235}
236
237#[cfg(feature = "python")]
238#[pymodule]
239fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
240    m.add_function(wrap_pyfunction!(generate, m)?)?;
241    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
242    m.add_function(wrap_pyfunction!(default_singleton_master_fraction, m)?)?;
243    m.add_class::<GenerateResult>()?;
244    Ok(())
245}
dupehell_core/lib.rs

dupehell_core/
lib.rs
