dupehell 0.6.0

// DupeHell -- MIT License
//
// Synthetic multi-domain dataset generator for record linkage benchmarking.
// No liability for misuse.

#[cfg(feature = "python")]
use pyo3::exceptions::PyValueError;
#[cfg(feature = "python")]
use pyo3::prelude::*;

mod buf_gen;
mod canary;
mod column_gen;
pub mod context;
pub mod cpu_affinity;
pub mod difficulty;
pub mod disk;
mod entity_gen;
mod fast_template;
mod fk_remap;
mod graph_gen;
pub mod gt;
mod hn_common;
mod noise;
pub mod pipeline;
mod pool_lookup;
pub mod rng;
pub mod schema;

pub use context::Context;
pub use difficulty::DifficultyReport;
pub use pipeline::{
    PipelineConfig, PipelineOutput, PipelineStats, run_pipeline, run_pipeline_with_progress,
};
pub use schema::{
    DomainSchema, EntitySchema, HnSchema, build_pipeline_config, chrono_now, load_schema,
};

// `GenerateResult` is only ever constructed in Rust and returned to Python
// (never passed back in as an argument). Since pyo3 0.30, `Clone` pyclasses
// no longer get a `FromPyObject` impl by default, which is what we want here.
#[cfg(feature = "python")]
#[pyclass]
#[derive(Clone)]
pub struct GenerateResult {
    #[pyo3(get)]
    pub dataset: String,
    #[pyo3(get)]
    pub ground_truth: String,
    #[pyo3(get)]
    pub total_records: usize,
    #[pyo3(get)]
    pub exact_dups: usize,
    #[pyo3(get)]
    pub fuzzy_dups: usize,
    #[pyo3(get)]
    pub hard_negs: usize,
    #[pyo3(get)]
    pub uniques: usize,
    #[pyo3(get)]
    pub masters: usize,
    #[pyo3(get)]
    pub nodes: Option<String>,
    #[pyo3(get)]
    pub edges: Option<String>,
}

#[cfg(feature = "python")]
#[pymethods]
impl GenerateResult {
    fn __repr__(&self) -> String {
        format!(
            "GenerateResult(dataset={:?}, ground_truth={:?}, total_records={}, exact_dups={}, fuzzy_dups={}, hard_negs={}, uniques={}, masters={})",
            self.dataset,
            self.ground_truth,
            self.total_records,
            self.exact_dups,
            self.fuzzy_dups,
            self.hard_negs,
            self.uniques,
            self.masters,
        )
    }

    fn __str__(&self) -> String {
        self.__repr__()
    }
}

/// The chosen difficulty tier's own singleton-master fraction (0.50/0.30/
/// 0.10 for light/medium/hell). Exposed so the Python `generate()`
/// wrapper can derive its default from `difficulty` instead of hardcoding a
/// single fixed value that silently ignores the tier — see
/// `schema::default_singleton_master_fraction` for the single source of
/// truth this mirrors.
#[cfg(feature = "python")]
#[pyfunction]
fn default_singleton_master_fraction(difficulty: &str) -> f64 {
    schema::default_singleton_master_fraction(difficulty)
}

#[cfg(feature = "python")]
#[pyfunction]
fn estimate_difficulty(
    domain: &str,
    size: usize,
    seed: u64,
    difficulty: &str,
    schemas_dir: &str,
    hard_neg_ratio: f64,
) -> PyResult<String> {
    let schema =
        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;
    let report = crate::difficulty::estimate_difficulty(
        domain,
        size,
        seed,
        difficulty,
        hard_neg_ratio,
        &schema,
    )
    .map_err(PyValueError::new_err)?;
    serde_json::to_string(&report).map_err(|e| PyValueError::new_err(e.to_string()))
}

#[cfg(feature = "python")]
#[pyfunction]
#[pyo3(signature = (
    domain, size, seed, difficulty, output_dir, locale, pools_dir, schemas_dir,
    output_format, hard_neg_ratio, singleton_master_fraction, generate_graph,
    graph_format, only_entity=None, chunk_size=None,
))]
#[allow(clippy::too_many_arguments)]
fn generate(
    domain: &str,
    size: usize,
    seed: u64,
    difficulty: &str,
    output_dir: &str,
    locale: &str,
    pools_dir: &str,
    schemas_dir: &str,
    output_format: &str,
    hard_neg_ratio: f64,
    singleton_master_fraction: f64,
    generate_graph: bool,
    graph_format: &str,
    only_entity: Option<&str>,
    chunk_size: Option<usize>,
) -> PyResult<GenerateResult> {
    let schema =
        load_schema(domain, std::path::Path::new(schemas_dir)).map_err(PyValueError::new_err)?;

    if output_format != "ipc" && output_format != "parquet" {
        return Err(PyValueError::new_err(format!(
            "invalid output format '{output_format}'; expected 'ipc' or 'parquet'"
        )));
    }
    // Pool files key their data by these exact lowercase codes; anything
    // else (e.g. "FR") would silently fall back to "en" in `Context::new`.
    if !crate::context::LOCALES.contains(&locale) {
        return Err(PyValueError::new_err(format!(
            "invalid locale '{locale}'; expected one of {:?}",
            crate::context::LOCALES
        )));
    }
    if !schema::valid_difficulty_names().contains(&difficulty) {
        return Err(PyValueError::new_err(format!(
            "invalid difficulty '{difficulty}'; expected one of {:?}",
            schema::valid_difficulty_names()
        )));
    }

    let mut ctx = Context::new(domain, locale, pools_dir).map_err(PyValueError::new_err)?;

    let run_id = schema::deterministic_run_id(
        domain,
        size,
        seed,
        difficulty,
        hard_neg_ratio,
        singleton_master_fraction,
        locale,
        only_entity,
        chunk_size,
    );

    let config = build_pipeline_config(
        domain,
        size,
        seed,
        difficulty,
        hard_neg_ratio,
        singleton_master_fraction,
        &schema,
        &run_id,
        output_format,
        generate_graph,
        graph_format,
        only_entity,
    )
    .map_err(PyValueError::new_err)?;
    crate::pipeline::source_date_epoch().map_err(PyValueError::new_err)?;
    crate::pipeline::check_capacity(&config).map_err(PyValueError::new_err)?;
    let chunked = matches!(chunk_size, Some(cs) if cs > 0 && cs < size);
    let needed = crate::pipeline::estimate_output_bytes(&config, chunked);
    if let Some(free) = crate::disk::available_space(std::path::Path::new(output_dir))
        && needed > free
    {
        return Err(PyValueError::new_err(format!(
            "this run needs ~{:.1} GB of disk space, but only {:.1} GB is free on the disk \
             holding {output_dir} — free some space, choose another output_dir, use a \
             smaller size, or output_format=\"parquet\" (6-12x smaller)",
            crate::disk::gb(needed),
            crate::disk::gb(free)
        )));
    }

    let output = match chunk_size {
        Some(cs) if cs > 0 && cs < size => {
            ctx.enable_watermark(domain, size, seed);
            crate::pipeline::run_chunked(
                &ctx,
                domain,
                size,
                cs,
                seed,
                difficulty,
                hard_neg_ratio,
                singleton_master_fraction,
                &schema,
                &run_id,
                output_format,
                generate_graph,
                graph_format,
                only_entity,
                output_dir,
                None,
                false, // skip_ground_truth: not yet exposed via the Python binding
            )
            .map_err(PyValueError::new_err)?
        }
        _ => {
            ctx.enable_watermark(&config.domain, config.size, config.seed);
            run_pipeline(&ctx, &config, output_dir).map_err(PyValueError::new_err)?
        }
    };

    Ok(GenerateResult {
        dataset: output.output_files.into_iter().next().unwrap_or_default(),
        ground_truth: output.gt_file,
        total_records: output.stats.total_records,
        exact_dups: output.stats.exact_dups,
        fuzzy_dups: output.stats.fuzzy_dups,
        hard_negs: output.stats.hard_negs,
        uniques: output.stats.uniques,
        masters: output.stats.masters,
        nodes: output.nodes,
        edges: output.edges,
    })
}

#[cfg(feature = "python")]
#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(generate, m)?)?;
    m.add_function(wrap_pyfunction!(estimate_difficulty, m)?)?;
    m.add_function(wrap_pyfunction!(default_singleton_master_fraction, m)?)?;
    m.add_class::<GenerateResult>()?;
    Ok(())
}