use std::collections::HashMap;
use std::sync::Arc;
use arrow::array::{
Array, ArrayBuilder, ArrayRef, AsArray, StringArray, UInt64Array, UInt64Builder,
};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use serde::Deserialize;
use crate::context::Context;
use crate::rng::Rng;
#[derive(Debug, Deserialize)]
pub struct PipelineConfig {
pub domain: String,
pub size: usize,
pub seed: u64,
pub difficulty: String,
pub output_format: String,
pub run_id: String,
pub entity_plans: Vec<EntityPlan>,
pub hard_neg_types: Vec<HnTypeConfig>,
pub hard_neg_ratio: f64,
pub graph_enabled: bool,
pub graph_format: String,
}
#[derive(Debug, Deserialize)]
pub struct EntityPlan {
pub name: String,
pub n_base: usize,
pub identifier_col: Option<String>,
pub columns_json: String,
pub noise_types: Vec<NoisePlanEntry>,
pub fk_remaps: Vec<FkRemapEntry>,
}
#[derive(Debug, Deserialize)]
pub struct NoisePlanEntry {
pub noise_type: String,
pub columns: Vec<String>,
pub count: usize,
}
#[derive(Debug, Deserialize)]
pub struct FkRemapEntry {
pub source_col: String,
pub target_entity: String,
}
#[derive(Debug, Deserialize)]
pub struct HnTypeConfig {
pub entity_type: String,
pub config_json: String,
pub count: usize,
}
#[derive(Debug)]
pub struct PipelineOutput {
pub output_files: Vec<String>,
pub gt_file: String,
pub stats: PipelineStats,
pub nodes: Option<String>,
pub edges: Option<String>,
}
#[derive(Debug)]
pub struct PipelineStats {
pub total_records: usize,
pub exact_dups: usize,
pub hard_negs: usize,
pub uniques: usize,
pub masters: usize,
}
const RID_LEN: usize = 15; const PAD_LEN: usize = 13;
#[inline]
fn record_id_string(i: usize) -> String {
let mut buf = [0u8; RID_LEN];
buf[0] = b'R';
buf[1] = b'-';
let mut n = i;
for j in (2..RID_LEN).rev() {
buf[j] = b'0' + (n % 10) as u8;
n /= 10;
}
String::from_utf8(buf.to_vec()).unwrap()
}
#[inline]
fn pad_string(i: usize) -> String {
let mut buf = [0u8; PAD_LEN];
let mut n = i;
for j in (0..PAD_LEN).rev() {
buf[j] = b'0' + (n % 10) as u8;
n /= 10;
}
String::from_utf8(buf.to_vec()).unwrap()
}
pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
range.map(record_id_string).collect()
}
fn entity_prefix(name: &str) -> String {
let h: u64 = name
.bytes()
.fold(0u64, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u64));
format!("E{:05}", h % 100000)
}
type FkPoolMap = HashMap<String, RecordBatch>;
struct HnPool {
batch: RecordBatch,
total_count: usize,
record_ids: Vec<String>,
}
const PERSON_NAME_WORDS: &[&str] = &[
"name",
"first",
"last",
"given",
"family",
"operator",
"technician",
"inspector",
"claimant",
"consignee",
"investigator",
"collected_by",
"performed_by",
"assigned_to",
];
const COMPANY_NAME_WORDS: &[&str] = &[
"company",
"legal",
"trading",
"name",
"supplier",
"manufacturer",
"sponsor",
"institution",
"journal",
"funder",
"employer",
"law_firm",
"affiliation",
"reported_by",
"buyer",
"shipper",
"airline",
];
pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
let lower = col_name.to_lowercase();
match noise_type {
"typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
| "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
| "blocking_fail" => {
if lower.contains("email") {
return false;
}
contains_any(&lower, &["address", "street", "city", "phone"])
|| contains_any(&lower, PERSON_NAME_WORDS)
|| contains_any(&lower, COMPANY_NAME_WORDS)
}
"names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
contains_any(&lower, PERSON_NAME_WORDS)
}
"dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
contains_any(&lower, &["date", "birth", "incorporation", "founding"])
}
"missing" | "missing_pattern" => {
contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
}
"identifiers"
| "corrupt_email"
| "corrupt_phone"
| "corrupt_national_id"
| "corrupt_siren"
| "national_id_corrupt"
| "phone_corrupt"
| "email_corrupt"
| "siren_corrupt" => contains_any(
&lower,
&[
"email",
"phone",
"ssn",
"pan",
"passport",
"account",
"medicare",
"license",
"siren",
"national_id",
],
),
"extra"
| "name_null"
| "dob_null"
| "blocking_fail_initial"
| "blocking_fail_partial"
| "fuzzy_match"
| "phonetic" => {
if lower.contains("email") {
return false;
}
contains_any(
&lower,
&["phone", "address", "street", "city", "note", "comment"],
) || contains_any(&lower, PERSON_NAME_WORDS)
|| contains_any(&lower, COMPANY_NAME_WORDS)
}
"companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
contains_any(&lower, COMPANY_NAME_WORDS)
}
"addresses" | "address_scramble" | "postal_corrupt" => {
contains_any(&lower, &["address", "street", "postal", "city"])
}
"exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
| "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
| "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
| "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
| "combo_ultimate" | "french_address" => {
false
}
_ => true,
}
}
fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
let mut matched: Vec<String> = Vec::new();
for field in schema.fields() {
if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
continue;
}
if exclude_cols.contains(field.name()) {
continue;
}
if noise_type_targets_column(noise_type, field.name()) {
matched.push(field.name().clone());
}
}
if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
matched.truncate(2);
}
matched.sort();
matched.dedup();
matched
}
fn contains_any(s: &str, patterns: &[&str]) -> bool {
patterns.iter().any(|p| s.contains(p))
}
fn apply_noise_to_batch(
rb: &RecordBatch,
noise_type: &str,
plan_cols: &[String],
rng: &mut Rng,
exclude_cols: &[String],
) -> Result<RecordBatch, String> {
let schema = rb.schema();
let n_cols = rb.num_columns();
let target_cols: Vec<String> = if plan_cols.is_empty() {
match_noise_columns(&schema, noise_type, exclude_cols)
} else {
plan_cols.to_vec()
};
let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();
let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
new_columns[idx]
.clone()
.unwrap_or_else(|| rb.column(idx).clone())
};
match noise_type {
"blocking_fail" => {
for col_name in &target_cols {
let lower = col_name.to_lowercase();
if let Ok(idx) = schema.index_of(col_name) {
let col = get_col(&new_columns, idx, rb);
if lower.contains("first") || lower.contains("given") {
new_columns[idx] =
Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
} else {
new_columns[idx] =
Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
}
}
}
}
"swap" | "full_swap" => {
if target_cols.len() >= 2 {
let name_a = &target_cols[0];
let name_b = &target_cols[1];
if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
let a = get_col(&new_columns, idx_a, rb);
let b = get_col(&new_columns, idx_b, rb);
new_columns[idx_a] = Some(b);
new_columns[idx_b] = Some(a);
}
}
}
"missing_pattern" => {
let mut phone_idx: Option<usize> = None;
let mut email_idx: Option<usize> = None;
for col_name in &target_cols {
let lower = col_name.to_lowercase();
if let Ok(idx) = schema.index_of(col_name) {
if lower.contains("phone") || lower.contains("mobile") {
phone_idx = Some(idx);
} else if lower.contains("email") {
email_idx = Some(idx);
}
}
}
if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
let phone_arr = get_col(&new_columns, pi, rb);
let email_arr = get_col(&new_columns, ei, rb);
let phone_s = phone_arr.as_string::<i32>();
let email_s = email_arr.as_string::<i32>();
let n = phone_s.len();
let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
for i in 0..n {
let pv = if phone_s.is_null(i) {
""
} else {
phone_s.value(i)
};
let ev = if email_s.is_null(i) {
""
} else {
email_s.value(i)
};
if !ev.is_empty() && rng.next_usize(2) == 0 {
builder.append_value("");
} else {
builder.append_value(pv);
}
}
new_columns[pi] = Some(Arc::new(builder.finish()));
}
}
_ => {
for col_name in &target_cols {
if let Ok(col_idx) = schema.index_of(col_name) {
let col = get_col(&new_columns, col_idx, rb);
let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
.map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
new_columns[col_idx] = Some(result);
}
}
}
}
let final_columns: Vec<ArrayRef> = (0..n_cols)
.map(|i| {
new_columns[i]
.clone()
.unwrap_or_else(|| rb.column(i).clone())
})
.collect();
RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))
}
fn topological_sort(
plans: &[EntityPlan],
name_to_idx: &HashMap<&str, usize>,
) -> Result<Vec<usize>, String> {
let n = plans.len();
let mut in_degree = vec![0usize; n];
let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
for i in 0..n {
for remap in &plans[i].fk_remaps {
let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
continue;
};
if target == i {
continue;
}
adj[target].push(i); in_degree[i] += 1;
}
}
let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
let mut order = Vec::with_capacity(n);
let mut idx = 0;
while idx < queue.len() {
let v = queue[idx];
idx += 1;
order.push(v);
for &next in &adj[v] {
in_degree[next] -= 1;
if in_degree[next] == 0 {
queue.push(next);
}
}
}
if order.len() != n {
return Err("FK dependency cycle detected among entities".into());
}
Ok(order)
}
pub fn run_pipeline(
ctx: &Context,
config: &PipelineConfig,
output_dir: &str,
) -> Result<PipelineOutput, String> {
std::fs::create_dir_all(output_dir)
.map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
let t_start = std::time::Instant::now();
let t_alloc = t_start.elapsed().as_secs_f64();
let t1e = std::time::Instant::now();
let metadata = build_metadata_map(config);
let full_arc = Arc::new(build_full_schema(config, &metadata));
let mut dataset_path = format!("{}/{}.ipc", output_dir, config.run_id);
let out_file =
std::fs::File::create(&dataset_path).map_err(|e| format!("create {dataset_path}: {e}"))?;
let mut writer = arrow::ipc::writer::FileWriter::try_new(out_file, &full_arc)
.map_err(|e| format!("IPC FileWriter error: {e}"))?;
let name_to_idx: HashMap<&str, usize> = config
.entity_plans
.iter()
.enumerate()
.map(|(i, p)| (p.name.as_str(), i))
.collect();
let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;
let mut fk_pools: FkPoolMap = HashMap::new();
let mut hn_pools: HashMap<String, HnPool> = HashMap::new();
let gt_ext = if config.output_format == "parquet" {
"parquet"
} else {
"ipc"
};
let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
let gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
let mut gt_acc = crate::gt::GtAccumulator::new(>_draft_path)?;
let mut global_rid_offset = 0usize;
const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
let mut graph_nodes_path: Option<String> = None;
let mut graph_edges_path: Option<String> = None;
if config.graph_enabled {
let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
node_writer = Some(
crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
.map_err(|e| format!("init node writer: {e}"))?,
);
edge_writer = Some(
crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata)
.map_err(|e| format!("init edge writer: {e}"))?,
);
graph_nodes_path = Some(nodes_ipc);
graph_edges_path = Some(edges_ipc);
}
let mut _t_write = 0.0f64;
let mut _t_meta = 0.0f64;
let mut _t_dup = 0.0f64;
let mut _write_calls = 0u64;
let hn_pool_sizing: std::collections::HashMap<&str, usize> =
config
.hard_neg_types
.iter()
.fold(std::collections::HashMap::new(), |mut m, hn| {
let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
*entry = (*entry).max(hn.count);
m
});
let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
const FK_POOL_CAP: usize = 200_000;
let fk_targets: std::collections::HashSet<&str> = config
.entity_plans
.iter()
.flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
.collect();
for &plan_idx in &entity_order {
let plan = &config.entity_plans[plan_idx];
if plan.n_base == 0 {
continue;
}
let prefix = entity_prefix(&plan.name);
let col_json_str = &plan.columns_json;
let fk_targeted = fk_targets.contains(plan.name.as_str());
let mut fk_builder = arrow::array::StringBuilder::new();
let mut fk_rid_builder = arrow::array::StringBuilder::new();
let mut fk_count: usize = 0;
let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
let mut hn_slices: Vec<RecordBatch> = Vec::new();
let mut hn_slice_rids: Vec<Vec<String>> = Vec::new();
let mut last_batch: Option<(RecordBatch, usize, usize)> = None;
let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
let mut fk_rng = Rng::new(config.seed.wrapping_add(42));
let mut col_lookup: Option<Vec<Option<usize>>> = None;
let mut null_cache: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
let mut const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();
for (batch_idx, offset) in (0..plan.n_base)
.step_by(crate::entity_gen::BATCH_SIZE)
.enumerate()
{
let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
let batch_seed = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);
let request_json = format!(
r#"{{"entity_name":"{}","n":{},"seed":{},"columns":{}}}"#,
plan.name, batch_n, batch_seed, col_json_str,
);
let rb = crate::entity_gen::generate_entity_batch(ctx, &request_json)?;
let mut fk_edges_by_remap: Vec<(String, Vec<String>)> = Vec::new();
let rb = if !plan.fk_remaps.is_empty() {
let mut r = rb;
for remap in &plan.fk_remaps {
if let Some(pool) = fk_pools.get(&remap.target_entity) {
let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
&r,
pool,
&remap.source_col,
&mut fk_rng,
config.graph_enabled,
)?;
r = remapped;
if config.graph_enabled
&& let Some(rids) = target_rids
{
fk_edges_by_remap.push((remap.source_col.clone(), rids));
}
}
}
r
} else {
rb
};
if fk_targeted
&& fk_count < FK_POOL_CAP
&& let Some(ref id_col) = plan.identifier_col
&& let Some(col) = rb.column_by_name(id_col)
{
let s = col.as_string::<i32>();
for i in 0..s.len() {
if fk_count >= FK_POOL_CAP {
break;
}
if !s.is_null(i) {
fk_builder.append_value(s.value(i));
if config.graph_enabled {
fk_rid_builder.append_value(record_id_string(global_rid_offset + i));
}
fk_count += 1;
}
}
}
let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
if hn_targeted && hn_accum < hn_max_pool {
let take = (hn_max_pool - hn_accum).min(batch_n);
hn_slices.push(rb.slice(0, take));
if config.graph_enabled {
hn_slice_rids.push(record_id_strs(global_rid_offset..global_rid_offset + take));
}
}
if col_lookup.is_none() {
let plan_schema = rb.schema();
col_lookup = Some(
full_arc
.fields()
.iter()
.skip(4)
.map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
.collect(),
);
}
let batch_mids: Vec<String> = (offset..offset + batch_n)
.map(|i| format!("{}-{}", prefix, pad_string(i)))
.collect();
let batch_mids_slice = &batch_mids[..];
let rid_slice = record_id_strs(global_rid_offset..global_rid_offset + batch_n);
let t_m0 = std::time::Instant::now();
let base_rb = add_metadata_and_align(
&rb,
&config.domain,
&plan.name,
&rid_slice,
batch_mids_slice,
&full_arc,
col_lookup.as_ref().unwrap(),
&mut null_cache,
&mut const_arr_cache,
);
_t_meta += t_m0.elapsed().as_secs_f64();
if config.graph_enabled
&& let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
{
nw.write_batch(&base_rb)
.map_err(|e| format!("write node: {e}"))?;
for i in 0..batch_n {
for (subtype, target_rids) in &fk_edges_by_remap {
ew.push(&rid_slice[i], &target_rids[i], "fk", subtype, 1.0)
.map_err(|e| format!("push fk edge: {e}"))?;
}
}
}
let t_w0 = std::time::Instant::now();
writer
.write(&base_rb)
.map_err(|e| format!("write base: {e}"))?;
_t_write += t_w0.elapsed().as_secs_f64();
_write_calls += 1;
gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;
global_rid_offset += batch_n;
last_batch = Some((rb, batch_n, offset));
}
if let Some(id_col) = &plan.identifier_col
&& fk_builder.len() > 0
{
let arr = Arc::new(fk_builder.finish()) as ArrayRef;
if config.graph_enabled {
let rid_arr = Arc::new(fk_rid_builder.finish()) as ArrayRef;
let schema = Arc::new(Schema::new(vec![
Field::new(id_col, DataType::Utf8, true),
Field::new("record_id", DataType::Utf8, true),
]));
if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
fk_pools.insert(plan.name.clone(), pool_rb);
}
} else {
let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
fk_pools.insert(plan.name.clone(), pool_rb);
}
}
}
if !hn_slices.is_empty() {
let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
let batch = if hn_slices.len() == 1 {
let only = hn_slices.into_iter().next().unwrap();
let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
pick_rows(&only, &idx)?
} else {
let schema = hn_slices[0].schema();
let n_fields = schema.fields().len();
let mut concat_arrays = Vec::with_capacity(n_fields);
for i in 0..n_fields {
let refs: Vec<&dyn Array> =
hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
concat_arrays.push(arrow::compute::concat(&refs).expect("hn pool concat col"));
}
RecordBatch::try_new(schema, concat_arrays).expect("hn pool concat batch")
};
let record_ids = if config.graph_enabled {
hn_slice_rids.into_iter().flatten().collect::<Vec<String>>()
} else {
Vec::new()
};
hn_pools.insert(
plan.name.clone(),
HnPool {
batch,
total_count,
record_ids,
},
);
}
let has_dups: bool = plan.noise_types.iter().any(|n| n.count > 0);
if has_dups && let Some((ref last_rb, last_n, last_offset)) = last_batch {
let t_d0 = std::time::Instant::now();
let mut dup_batches: Vec<RecordBatch> = Vec::new();
let mut dup_mids_buf: Vec<String> = Vec::new();
let fk_exclude_cols: Vec<String> = plan
.fk_remaps
.iter()
.map(|r| r.source_col.clone())
.collect();
let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
.noise_types
.iter()
.filter(|n| n.count > 0)
.map(|n| {
let mut b = UInt64Builder::with_capacity(n.count);
for _ in 0..n.count {
b.append_value(batch_rng.next_usize(last_n) as u64);
}
(
b.finish(),
batch_rng.next_u64(),
n.noise_type.as_str(),
n.columns.as_slice(),
n.count,
)
})
.collect();
let mut results: Vec<Result<(RecordBatch, Vec<String>), String>> =
Vec::with_capacity(ndata.len());
std::thread::scope(|s| {
let mut handles = Vec::with_capacity(ndata.len());
for (indices, seed, ntype, cols, cnt) in &ndata {
let rb = last_rb.clone();
let idxs = indices.clone();
let cols_v: Vec<String> = cols.to_vec();
let exclude = fk_exclude_cols.clone();
let prefix_ref: &str = &prefix;
handles.push(s.spawn(move || {
let dup = pick_rows(&rb, &idxs)?;
let mut rng = Rng::new(*seed);
let noisy = apply_noise_to_batch(&dup, ntype, &cols_v, &mut rng, &exclude)?;
let mut mb = Vec::with_capacity(*cnt);
for j in 0..*cnt {
let global_idx = last_offset + idxs.value(j) as usize;
mb.push(format!("{}-{}", prefix_ref, pad_string(global_idx)));
}
Ok((noisy, mb))
}));
}
for h in handles {
results.push(h.join().unwrap());
}
});
for res in results {
let (rb, mb) = res?;
dup_batches.push(rb);
dup_mids_buf.extend(mb);
}
_t_dup += t_d0.elapsed().as_secs_f64();
if !dup_batches.is_empty() {
let dup_total = dup_mids_buf.len();
let dup_rids = record_id_strs(global_rid_offset..global_rid_offset + dup_total);
let concated = if dup_batches.len() == 1 {
dup_batches.into_iter().next().unwrap()
} else {
let schema = dup_batches[0].schema();
let n_fields = dup_batches[0].num_columns();
let mut concat_arrays = Vec::with_capacity(n_fields);
for i in 0..n_fields {
let refs: Vec<&dyn Array> =
dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
concat_arrays.push(
arrow::compute::concat(&refs)
.map_err(|e| format!("concat dup col {i}: {e}"))?,
);
}
RecordBatch::try_new(schema, concat_arrays)
.map_err(|e| format!("concat dups: {e}"))?
};
let dup_rb_full = add_metadata_and_align(
&concated,
&config.domain,
&plan.name,
&dup_rids,
&dup_mids_buf,
&full_arc,
col_lookup.as_ref().unwrap(),
&mut null_cache,
&mut const_arr_cache,
);
let t_wd = std::time::Instant::now();
writer
.write(&dup_rb_full)
.map_err(|e| format!("write dups: {e}"))?;
_t_write += t_wd.elapsed().as_secs_f64();
_write_calls += 1;
if config.graph_enabled
&& let Some(nw) = node_writer.as_mut()
{
nw.write_batch(&dup_rb_full)
.map_err(|e| format!("write dup node: {e}"))?;
}
gt_acc.push_dup_batch(
dup_rb_full.column(0),
dup_rb_full.column(2),
dup_rb_full.column(3),
)?;
global_rid_offset += dup_total;
}
}
}
let t1e_elapsed = t1e.elapsed().as_secs_f64();
log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
let t2 = std::time::Instant::now();
let mut hn_null_cache: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
let mut hn_const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();
for (hn_idx, hn_cfg) in config.hard_neg_types.iter().enumerate() {
if hn_cfg.count == 0 {
continue;
}
let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
continue;
};
let hn_total = pool_data.total_count;
let pool_rb = &pool_data.batch;
let n_hn_max = hn_cfg.count.min(hn_total / 2);
if n_hn_max == 0 {
continue;
}
let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
pool_rb,
&hn_cfg.config_json,
n_hn_max,
config.seed.wrapping_add(200),
config.graph_enabled,
)?;
let n_hn = hn_rb.num_rows();
if n_hn == 0 {
continue;
}
let hn_rids = record_id_strs(global_rid_offset..global_rid_offset + n_hn);
let mut hn_mid_rng = Rng::new(config.seed.wrapping_add(300 + hn_idx as u64));
let hn_mids: Vec<String> = (0..n_hn)
.map(|_i| {
let pad = hn_mid_rng.next_usize(10_000_000);
format!("HN-{:07}", pad)
})
.collect();
let hn_schema = hn_rb.schema();
let hn_col_lookup: Vec<Option<usize>> = full_arc
.fields()
.iter()
.skip(4)
.map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
.collect();
let hn_rb_full = add_metadata_and_align(
&hn_rb,
&config.domain,
&hn_cfg.entity_type,
&hn_rids,
&hn_mids,
&full_arc,
&hn_col_lookup,
&mut hn_null_cache,
&mut hn_const_arr_cache,
);
let t_wh = std::time::Instant::now();
writer
.write(&hn_rb_full)
.map_err(|e| format!("write hn: {e}"))?;
_t_write += t_wh.elapsed().as_secs_f64();
_write_calls += 1;
if config.graph_enabled
&& let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
{
nw.write_batch(&hn_rb_full)
.map_err(|e| format!("write hn node: {e}"))?;
if let Some((idx_a, pattern)) = hn_src {
for i in 0..n_hn {
let tgt = &pool_data.record_ids[idx_a[i]];
ew.push(&hn_rids[i], tgt, "hard_neg", &pattern, 1.0)
.map_err(|e| format!("push hn edge: {e}"))?;
}
}
}
gt_acc.push_other_batch(
hn_rb_full.column(0),
hn_rb_full.column(2),
hn_rb_full.column(3),
)?;
global_rid_offset += n_hn;
}
let t2_elapsed = t2.elapsed().as_secs_f64();
{
let mut canary_null_cache: HashMap<(DataType, usize), ArrayRef> = HashMap::new();
let mut canary_const_arr_cache: HashMap<(String, usize), ArrayRef> = HashMap::new();
crate::canary::generate_all(
ctx,
config,
&full_arc,
&mut canary_null_cache,
&mut canary_const_arr_cache,
&mut global_rid_offset,
&fk_pools,
&mut writer,
&mut node_writer,
&mut gt_acc,
)?;
}
let t3 = std::time::Instant::now();
writer
.finish()
.map_err(|e| format!("finish IPC writer: {e}"))?;
let t3a_elapsed = t3.elapsed().as_secs_f64();
let t3b = std::time::Instant::now();
let t_gt0 = std::time::Instant::now();
let crate::gt::GtResult {
n_exact_dup,
n_hard_neg,
n_unique,
n_masters,
cluster_map,
} = gt_acc.finish(
config.difficulty.as_str(),
&config.output_format,
>_path,
&metadata,
)?;
let _t_gt_compute = 0.0f64;
let _t_gt_write = t_gt0.elapsed().as_secs_f64();
let t3b_elapsed = t3b.elapsed().as_secs_f64();
let _t3_elapsed = t3.elapsed().as_secs_f64();
if config.graph_enabled
&& let Some(ew) = edge_writer.as_mut()
{
crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
.map_err(|e| format!("push dup clusters: {e}"))?;
}
if config.output_format == "parquet" {
use arrow::ipc::reader::FileReader;
let parquet_path = dataset_path.replace(".ipc", ".parquet");
let ipc_file = std::fs::File::open(&dataset_path)
.map_err(|e| format!("open IPC for conversion: {e}"))?;
let ipc_reader = FileReader::try_new(ipc_file, None)
.map_err(|e| format!("IPC reader for conversion: {e}"))?;
let schema = ipc_reader.schema();
let parquet_file = std::fs::File::create(&parquet_path)
.map_err(|e| format!("create {parquet_path}: {e}"))?;
let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
.clone()
.into_iter()
.map(|(k, v)| parquet::file::metadata::KeyValue {
key: k,
value: Some(v),
})
.collect();
let props = parquet::file::properties::WriterProperties::builder()
.set_compression(parquet::basic::Compression::ZSTD(zstd_level))
.set_max_row_group_row_count(Some(1_000_000))
.set_key_value_metadata(Some(meta_kv))
.build();
let mut parquet_writer =
parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
.map_err(|e| format!("ArrowWriter for conversion: {e}"))?;
for batch_result in ipc_reader {
let batch = batch_result.map_err(|e| format!("IPC batch for conversion: {e}"))?;
parquet_writer
.write(&batch)
.map_err(|e| format!("write parquet conversion: {e}"))?;
}
parquet_writer
.close()
.map_err(|e| format!("close parquet conversion: {e}"))?;
std::fs::remove_file(&dataset_path).map_err(|e| format!("remove IPC temp file: {e}"))?;
dataset_path = parquet_path;
}
let duration = t_start.elapsed().as_secs_f64();
log::debug!(
"[pipeline_timing] alloc={:.3}s sink={:.3}s hn={:.3}s merge={:.3}s meta={:.3}s/dup={:.3}s/write={:.3}s gt(comp={:.3}s+write={:.3}s)={:.3}s total={:.3}s",
t_alloc,
t1e_elapsed,
t2_elapsed,
t3a_elapsed,
_t_meta,
_t_dup,
_t_write,
_t_gt_compute,
_t_gt_write,
t3b_elapsed,
duration,
);
let stats = PipelineStats {
total_records: global_rid_offset,
exact_dups: n_exact_dup,
hard_negs: n_hard_neg,
uniques: n_unique,
masters: n_masters,
};
let mut graph_nodes_final: Option<String> = None;
let mut graph_edges_final: Option<String> = None;
if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
nw.finish()
.map_err(|e| format!("finish node writer: {e}"))?;
ew.finish()
.map_err(|e| format!("finish edge writer: {e}"))?;
let mut np = graph_nodes_path.clone().unwrap();
let mut ep = graph_edges_path.clone().unwrap();
if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
std::fs::remove_file(&np).ok();
std::fs::remove_file(&ep).ok();
np = np_pq;
ep = ep_pq;
}
graph_nodes_final = Some(np);
graph_edges_final = Some(ep);
}
Ok(PipelineOutput {
output_files: vec![dataset_path],
gt_file: gt_path,
stats,
nodes: graph_nodes_final,
edges: graph_edges_final,
})
}
fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
use arrow::ipc::reader::FileReader;
let ipc_file =
std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
let ipc_reader =
FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
let schema = ipc_reader.schema();
let parquet_file =
std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
.metadata()
.iter()
.map(|(k, v)| parquet::file::metadata::KeyValue {
key: k.clone(),
value: Some(v.clone()),
})
.collect();
let props = parquet::file::properties::WriterProperties::builder()
.set_compression(parquet::basic::Compression::ZSTD(zstd_level))
.set_max_row_group_row_count(Some(1_000_000))
.set_key_value_metadata(Some(meta_kv))
.build();
let mut parquet_writer =
parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
.map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;
for batch_result in ipc_reader {
let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
parquet_writer
.write(&batch)
.map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
}
parquet_writer
.close()
.map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
Ok(())
}
fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs().to_string())
.unwrap_or_else(|_| "unknown".into());
HashMap::from([
(
"dupehell.generator".into(),
format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
),
(
"dupehell.provenance".into(),
"dupehell-synthetic-data".into(),
),
("dupehell.license".into(), "MIT".into()),
(
"dupehell.purpose".into(),
"Educational Use Only — Record Linkage Benchmarking".into(),
),
(
"dupehell.url".into(),
"https://github.com/vntoinekaio/DupeHell".into(),
),
("dupehell.domain".into(), config.domain.clone()),
("dupehell.size".into(), config.size.to_string()),
("dupehell.seed".into(), config.seed.to_string()),
("dupehell.difficulty".into(), config.difficulty.clone()),
("dupehell.run_id".into(), config.run_id.clone()),
("dupehell.timestamp".into(), ts),
(
"dupehell.structural_key_columns".into(),
structural_key_columns_note(config),
),
])
}
fn structural_key_columns_note(config: &PipelineConfig) -> String {
let cols: Vec<&str> = config
.entity_plans
.iter()
.filter_map(|p| p.identifier_col.as_deref())
.collect();
format!(
"{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
cols.join(", ")
)
}
fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
let mut field_map: Vec<(String, DataType, bool)> = Vec::new();
let metadata_fields = ["record_id", "domain", "entity_type", "master_id"];
for mf in &metadata_fields {
field_map.push((mf.to_string(), DataType::Utf8, false));
}
for plan in &config.entity_plans {
let cols: Vec<serde_json::Value> =
serde_json::from_str(&plan.columns_json).unwrap_or_default();
for col in &cols {
let name = col["name"].as_str().unwrap_or("").to_string();
if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
continue;
}
let col_type = col_type_from_request(col);
field_map.push((name, col_type, true));
}
}
let fields: Vec<Field> = field_map
.into_iter()
.map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
.collect();
Schema::new(fields).with_metadata(metadata.clone())
}
fn col_type_from_request(v: &serde_json::Value) -> DataType {
match v["type"].as_str() {
Some("int") => DataType::Int64,
Some("float") => DataType::Float64,
Some("boolean") => DataType::Boolean,
Some("date") | Some("datetime") => DataType::Utf8,
_ => DataType::Utf8,
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn add_metadata_and_align(
rb: &RecordBatch,
domain: &str,
entity_type: &str,
record_ids: &[String],
master_ids: &[String],
full_arc: &Arc<Schema>,
col_lookup: &[Option<usize>],
null_cache: &mut HashMap<(DataType, usize), ArrayRef>,
const_arr_cache: &mut HashMap<(String, usize), ArrayRef>,
) -> RecordBatch {
let n = rb.num_rows();
let domain_arr = const_arr_cache
.entry((domain.to_string(), n))
.or_insert_with(|| {
Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
domain, n,
))) as ArrayRef
})
.clone();
let et_arr = const_arr_cache
.entry((entity_type.to_string(), n))
.or_insert_with(|| {
Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
entity_type,
n,
))) as ArrayRef
})
.clone();
let rid_arr = Arc::new(StringArray::from_iter_values(
record_ids.iter().map(|s| s.as_str()),
)) as ArrayRef;
let mid_arr = Arc::new(StringArray::from_iter_values(
master_ids.iter().map(|s| s.as_str()),
)) as ArrayRef;
let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
all_arrays.push(rid_arr);
all_arrays.push(domain_arr);
all_arrays.push(et_arr);
all_arrays.push(mid_arr);
for (i, maybe_idx) in col_lookup.iter().enumerate() {
match maybe_idx {
Some(idx) => all_arrays.push(rb.column(*idx).clone()),
None => {
let dt = full_arc.field(i + 4).data_type();
let arr = null_cache
.entry((dt.clone(), n))
.or_insert_with(|| arrow::array::new_null_array(dt, n));
all_arrays.push(Arc::clone(arr));
}
}
}
RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
.expect("add_metadata_and_align RecordBatch")
}
fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
use arrow::compute::take;
let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
.map(|i| take(rb.column(i), indices, None))
.collect::<Result<Vec<_>, _>>()
.map_err(|e| format!("take error: {e}"))?;
RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
}