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;
17pub mod disk;
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). Since pyo3 0.30, `Clone` pyclasses
41// no longer get a `FromPyObject` impl by default, which is what we want here.
42#[cfg(feature = "python")]
43#[pyclass]
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.10 for light/medium/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, chunk_size=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    chunk_size: Option<usize>,
150) -> PyResult<GenerateResult> {
151    let schema =
152        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
153
154    if output_format != "ipc" && output_format != "parquet" {
155        return Err(PyValueError::new_err(format!(
156            "invalid output format '{output_format}'; expected 'ipc' or 'parquet'"
157        )));
158    }
159    // Pool files key their data by these exact lowercase codes; anything
160    // else (e.g. "FR") would silently fall back to "en" in `Context::new`.
161    if !crate::context::LOCALES.contains(&locale) {
162        return Err(PyValueError::new_err(format!(
163            "invalid locale '{locale}'; expected one of {:?}",
164            crate::context::LOCALES
165        )));
166    }
167    if !schema::valid_difficulty_names().contains(&difficulty) {
168        return Err(PyValueError::new_err(format!(
169            "invalid difficulty '{difficulty}'; expected one of {:?}",
170            schema::valid_difficulty_names()
171        )));
172    }
173
174    let mut ctx = Context::new(domain, locale, pools_dir).map_err(PyValueError::new_err)?;
175
176    let run_id = schema::deterministic_run_id(
177        domain,
178        size,
179        seed,
180        difficulty,
181        hard_neg_ratio,
182        singleton_master_fraction,
183        locale,
184        only_entity,
185        chunk_size,
186    );
187
188    let config = build_pipeline_config(
189        domain,
190        size,
191        seed,
192        difficulty,
193        hard_neg_ratio,
194        singleton_master_fraction,
195        &schema,
196        &run_id,
197        output_format,
198        generate_graph,
199        graph_format,
200        only_entity,
201    )
202    .map_err(PyValueError::new_err)?;
203    crate::pipeline::source_date_epoch().map_err(PyValueError::new_err)?;
204    crate::pipeline::check_capacity(&config).map_err(PyValueError::new_err)?;
205    let chunked = matches!(chunk_size, Some(cs) if cs > 0 && cs < size);
206    let needed = crate::pipeline::estimate_output_bytes(&config, chunked);
207    if let Some(free) = crate::disk::available_space(std::path::Path::new(output_dir))
208        && needed > free
209    {
210        return Err(PyValueError::new_err(format!(
211            "this run needs ~{:.1} GB of disk space, but only {:.1} GB is free on the disk \
212             holding {output_dir} — free some space, choose another output_dir, use a \
213             smaller size, or output_format=\"parquet\" (6-12x smaller)",
214            crate::disk::gb(needed),
215            crate::disk::gb(free)
216        )));
217    }
218
219    let output = match chunk_size {
220        Some(cs) if cs > 0 && cs < size => {
221            ctx.enable_watermark(domain, size, seed);
222            crate::pipeline::run_chunked(
223                &ctx,
224                domain,
225                size,
226                cs,
227                seed,
228                difficulty,
229                hard_neg_ratio,
230                singleton_master_fraction,
231                &schema,
232                &run_id,
233                output_format,
234                generate_graph,
235                graph_format,
236                only_entity,
237                output_dir,
238                None,
239                false, // skip_ground_truth: not yet exposed via the Python binding
240            )
241            .map_err(PyValueError::new_err)?
242        }
243        _ => {
244            ctx.enable_watermark(&config.domain, config.size, config.seed);
245            run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?
246        }
247    };
248
249    Ok(GenerateResult {
250        dataset: output.output_files.into_iter().next().unwrap_or_default(),
251        ground_truth: output.gt_file,
252        total_records: output.stats.total_records,
253        exact_dups: output.stats.exact_dups,
254        fuzzy_dups: output.stats.fuzzy_dups,
255        hard_negs: output.stats.hard_negs,
256        uniques: output.stats.uniques,
257        masters: output.stats.masters,
258        nodes: output.nodes,
259        edges: output.edges,
260    })
261}
262
263#[cfg(feature = "python")]
264#[pymodule]
265fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
266    m.add_function(wrap_pyfunction!(generate, m)?)?;
267    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
268    m.add_function(wrap_pyfunction!(default_singleton_master_fraction, m)?)?;
269    m.add_class::<GenerateResult>()?;
270    Ok(())
271}
dupehell_core/lib.rs

dupehell_core/
lib.rs
