1use std::collections::HashMap;
8use std::sync::Arc;
9
10use arrow::array::{
11 Array, ArrayBuilder, ArrayRef, AsArray, StringArray, UInt64Array, UInt64Builder,
12};
13use arrow::datatypes::{DataType, Field, Schema};
14use arrow::record_batch::RecordBatch;
15use rayon::prelude::*;
16use rustc_hash::FxHashMap;
17use serde::Deserialize;
18
19use crate::context::Context;
20use crate::rng::Rng;
21
22#[derive(Debug, Deserialize)]
25pub struct PipelineConfig {
26 pub domain: String,
27 pub size: usize,
28 pub seed: u64,
29 pub difficulty: String,
30 pub output_format: String,
31 pub run_id: String,
32
33 pub entity_plans: Vec<EntityPlan>,
34 pub hard_neg_types: Vec<HnTypeConfig>,
35
36 pub hard_neg_ratio: f64,
37
38 pub noise_passes: usize,
43
44 pub graph_enabled: bool,
45 pub graph_format: String,
46}
47
48#[derive(Debug, Deserialize)]
49pub struct EntityPlan {
50 pub name: String,
51 pub n_base: usize,
52 pub identifier_col: Option<String>,
53 pub columns_json: String,
54 pub noise_types: Vec<NoisePlanEntry>,
55 pub fk_remaps: Vec<FkRemapEntry>,
56 #[serde(default)]
63 pub pool_only: bool,
64}
65
66pub(crate) const FK_POOL_CAP: usize = 200_000;
73
74#[derive(Debug, Deserialize)]
75pub struct NoisePlanEntry {
76 pub noise_type: String,
77 pub columns: Vec<String>,
78 pub count: usize,
79}
80
81#[derive(Debug, Deserialize)]
82pub struct FkRemapEntry {
83 pub source_col: String,
84 pub target_entity: String,
85 #[serde(default = "default_fk_weight")]
89 pub fk_weight: f64,
90}
91
92fn default_fk_weight() -> f64 {
93 1.0
94}
95
96#[derive(Debug, Deserialize)]
97pub struct HnTypeConfig {
98 pub entity_type: String,
99 pub config_json: String,
100 pub count: usize,
101}
102
103#[derive(Debug)]
104pub struct PipelineOutput {
105 pub output_files: Vec<String>,
106 pub gt_file: String,
107 pub stats: PipelineStats,
108 pub nodes: Option<String>,
109 pub edges: Option<String>,
110}
111
112#[derive(Debug)]
113pub struct PipelineStats {
114 pub total_records: usize,
115 pub exact_dups: usize,
118 pub fuzzy_dups: usize,
121 pub hard_negs: usize,
122 pub uniques: usize,
123 pub masters: usize,
124}
125
126const RID_LEN: usize = 15; const PAD_LEN: usize = 13; #[inline]
141pub(crate) fn record_id_string(i: usize) -> String {
142 let mut buf = [0u8; RID_LEN];
143 buf[0] = b'R';
144 buf[1] = b'-';
145 let mut n = i;
146 for j in (2..RID_LEN).rev() {
147 buf[j] = b'0' + (n % 10) as u8;
148 n /= 10;
149 }
150 String::from_utf8(buf.to_vec()).expect("record_id_string: buffer is 'R', '-' and ASCII digits")
151}
152
153#[inline]
154pub(crate) fn pad_string(i: usize) -> String {
155 let mut buf = [0u8; PAD_LEN];
156 let mut n = i;
157 for j in (0..PAD_LEN).rev() {
158 buf[j] = b'0' + (n % 10) as u8;
159 n /= 10;
160 }
161 String::from_utf8(buf.to_vec()).expect("pad_string: buffer is always ASCII digits")
162}
163
164pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
165 range.map(record_id_string).collect()
166}
167
168pub(crate) fn parse_record_idx(rid: &str) -> Option<u64> {
174 if rid.len() != RID_LEN || !rid.starts_with("R-") {
175 return None;
176 }
177 rid[2..].parse().ok()
178}
179
180pub(crate) fn pack_master_key(mid: &str) -> Option<u64> {
194 let bytes = mid.as_bytes();
195 if bytes.len() != 1 + 5 + 1 + PAD_LEN || bytes[0] != b'E' || bytes[6] != b'-' {
196 return None;
197 }
198 let entity_idx: u64 = mid[1..6].parse().ok()?;
199 let local_idx: u64 = mid[7..7 + PAD_LEN].parse().ok()?;
200 Some((entity_idx << 44) | local_idx)
201}
202
203fn distribute_by_weight(total: usize, weights: &[usize]) -> Vec<usize> {
211 let weight_sum: usize = weights.iter().sum();
212 if weight_sum == 0 || total == 0 {
213 return vec![0; weights.len()];
214 }
215 let raw: Vec<f64> = weights
216 .iter()
217 .map(|&w| total as f64 * w as f64 / weight_sum as f64)
218 .collect();
219 let mut counts: Vec<usize> = raw.iter().map(|&r| r as usize).collect();
220 let assigned: usize = counts.iter().sum();
221 let mut remainder = total - assigned;
222 let mut order: Vec<usize> = (0..counts.len()).collect();
223 order.sort_by(|&a, &b| {
224 let fa = raw[a] - counts[a] as f64;
225 let fb = raw[b] - counts[b] as f64;
226 fb.partial_cmp(&fa).unwrap_or(std::cmp::Ordering::Equal)
227 });
228 for &i in &order {
229 if remainder == 0 {
230 break;
231 }
232 counts[i] += 1;
233 remainder -= 1;
234 }
235 counts
236}
237
238pub(crate) fn entity_prefix(entity_idx: usize) -> String {
247 format!("E{:05}", entity_idx)
248}
249
250type FkPoolMap = HashMap<String, RecordBatch>;
253
254struct HnPool {
255 batch: RecordBatch,
256 total_count: usize,
257 record_ids: Vec<String>,
260}
261
262const PERSON_NAME_WORDS: &[&str] = &[
269 "name",
270 "first_name",
278 "last_name",
279 "given",
280 "family",
281 "middle",
282 "operator",
283 "technician",
284 "inspector",
285 "claimant",
286 "consignee",
287 "investigator",
288 "collected_by",
289 "performed_by",
290 "assigned_to",
291];
292
293const COMPANY_NAME_WORDS: &[&str] = &[
299 "company",
300 "legal",
301 "trading",
302 "name",
303 "supplier",
304 "manufacturer",
305 "sponsor",
306 "institution",
307 "journal",
308 "funder",
309 "employer",
310 "law_firm",
311 "affiliation",
312 "reported_by",
313 "buyer",
314 "shipper",
315 "airline",
316];
317
318fn is_specifically_person_name(lower: &str) -> bool {
332 PERSON_NAME_WORDS
333 .iter()
334 .any(|&w| w != "name" && lower.contains(w))
335}
336
337fn is_specifically_company_name(lower: &str) -> bool {
340 COMPANY_NAME_WORDS
341 .iter()
342 .any(|&w| w != "name" && lower.contains(w))
343}
344
345pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
351 let lower = col_name.to_lowercase();
352 match noise_type {
353 "typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
354 | "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
355 | "blocking_fail" => {
356 if lower.contains("email") || lower.contains("ip_address") {
361 return false;
362 }
363 contains_any(&lower, &["address", "street", "city"])
368 || contains_any(&lower, PERSON_NAME_WORDS)
369 || contains_any(&lower, COMPANY_NAME_WORDS)
370 }
371 "names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
372 contains_any(&lower, PERSON_NAME_WORDS) && !is_specifically_company_name(&lower)
378 }
379 "dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
380 contains_any(&lower, &["date", "birth", "incorporation", "founding"])
381 }
382 "missing" | "missing_pattern" => {
383 contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
384 }
385 "identifiers"
386 | "corrupt_email"
387 | "corrupt_phone"
388 | "corrupt_national_id"
389 | "corrupt_siren"
390 | "national_id_corrupt"
391 | "phone_corrupt"
392 | "email_corrupt"
393 | "siren_corrupt" => contains_any(
394 &lower,
395 &[
396 "email",
397 "phone",
398 "ssn",
399 "pan",
400 "passport",
401 "account",
402 "medicare",
403 "license",
404 "siren",
405 "national_id",
406 "registration",
413 ],
414 ),
415 "extra"
416 | "name_null"
417 | "dob_null"
418 | "blocking_fail_initial"
419 | "blocking_fail_partial"
420 | "fuzzy_match"
421 | "phonetic" => {
422 if lower.contains("email") || lower.contains("ip_address") {
428 return false;
429 }
430 contains_any(&lower, &["address", "street", "city", "note", "comment"])
431 || contains_any(&lower, PERSON_NAME_WORDS)
432 || contains_any(&lower, COMPANY_NAME_WORDS)
433 }
434 "companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
435 contains_any(&lower, COMPANY_NAME_WORDS) && !is_specifically_person_name(&lower)
441 }
442 "addresses" | "address_scramble" | "postal_corrupt" => {
443 !lower.contains("ip_address")
459 && !lower.contains("email")
460 && contains_any(
461 &lower,
462 &["address", "street", "postal", "city", "state", "country"],
463 )
464 }
465 "exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
466 | "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
467 | "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
468 | "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
469 | "combo_ultimate" | "french_address" => {
470 false
472 }
473 _ => true,
475 }
476}
477
478fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
479 let mut matched: Vec<String> = Vec::new();
480 for field in schema.fields() {
481 if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
482 continue;
483 }
484 if exclude_cols.contains(field.name()) {
485 continue;
486 }
487 if noise_type_targets_column(noise_type, field.name()) {
488 matched.push(field.name().clone());
489 }
490 }
491 if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
493 matched.truncate(2);
494 }
495 matched.sort();
496 matched.dedup();
497 matched
498}
499
500fn contains_any(s: &str, patterns: &[&str]) -> bool {
501 patterns.iter().any(|p| s.contains(p))
502}
503
504fn apply_noise_to_batch(
511 rb: &RecordBatch,
512 noise_type: &str,
513 plan_cols: &[String],
514 rng: &mut Rng,
515 exclude_cols: &[String],
516) -> Result<(RecordBatch, Vec<String>), String> {
517 let schema = rb.schema();
518 let n_cols = rb.num_columns();
519
520 let target_cols: Vec<String> = if plan_cols.is_empty() {
522 match_noise_columns(&schema, noise_type, exclude_cols)
523 } else {
524 plan_cols.to_vec()
525 };
526
527 let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();
529
530 let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
531 new_columns[idx]
532 .clone()
533 .unwrap_or_else(|| rb.column(idx).clone())
534 };
535
536 match noise_type {
537 "blocking_fail" => {
538 for col_name in &target_cols {
539 let lower = col_name.to_lowercase();
540 if let Ok(idx) = schema.index_of(col_name) {
541 let col = get_col(&new_columns, idx, rb);
542 if lower.contains("first") || lower.contains("given") {
543 new_columns[idx] =
544 Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
545 } else {
546 new_columns[idx] =
547 Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
548 }
549 }
550 }
551 }
552 "swap" | "full_swap" => {
553 if target_cols.len() >= 2 {
554 let name_a = &target_cols[0];
555 let name_b = &target_cols[1];
556 if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
557 let a = get_col(&new_columns, idx_a, rb);
558 let b = get_col(&new_columns, idx_b, rb);
559 new_columns[idx_a] = Some(b);
560 new_columns[idx_b] = Some(a);
561 }
562 }
563 }
564 "missing_pattern" => {
565 let mut phone_idx: Option<usize> = None;
566 let mut email_idx: Option<usize> = None;
567 for col_name in &target_cols {
568 let lower = col_name.to_lowercase();
569 if let Ok(idx) = schema.index_of(col_name) {
570 if lower.contains("phone") || lower.contains("mobile") {
571 phone_idx = Some(idx);
572 } else if lower.contains("email") {
573 email_idx = Some(idx);
574 }
575 }
576 }
577 if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
578 let phone_arr = get_col(&new_columns, pi, rb);
579 let email_arr = get_col(&new_columns, ei, rb);
580 let phone_s = phone_arr.as_string::<i32>();
581 let email_s = email_arr.as_string::<i32>();
582 let n = phone_s.len();
583 let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
584 for i in 0..n {
585 let pv = if phone_s.is_null(i) {
586 ""
587 } else {
588 phone_s.value(i)
589 };
590 let ev = if email_s.is_null(i) {
591 ""
592 } else {
593 email_s.value(i)
594 };
595 if !ev.is_empty() && rng.next_usize(2) == 0 {
596 builder.append_value("");
597 } else {
598 builder.append_value(pv);
599 }
600 }
601 new_columns[pi] = Some(Arc::new(builder.finish()));
602 }
603 }
604 _ => {
605 for col_name in &target_cols {
606 if let Ok(col_idx) = schema.index_of(col_name) {
607 let col = get_col(&new_columns, col_idx, rb);
608 let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
609 .map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
610 new_columns[col_idx] = Some(result);
611 }
612 }
613 }
614 }
615
616 let final_columns: Vec<ArrayRef> = (0..n_cols)
618 .map(|i| {
619 new_columns[i]
620 .clone()
621 .unwrap_or_else(|| rb.column(i).clone())
622 })
623 .collect();
624
625 let batch =
626 RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))?;
627 Ok((batch, target_cols))
628}
629
630fn unchanged_row_mask(
641 orig: &RecordBatch,
642 noised: &RecordBatch,
643 target_cols: &[String],
644) -> Result<Vec<bool>, String> {
645 let n = orig.num_rows();
646 let mut changed = vec![false; n];
647 let schema = orig.schema();
648 for col_name in target_cols {
649 let Ok(idx) = schema.index_of(col_name) else {
650 continue;
651 };
652 let a = orig.column(idx).as_string::<i32>();
653 let b = noised.column(idx).as_string::<i32>();
654 for (i, is_changed) in changed.iter_mut().enumerate() {
655 if *is_changed {
656 continue;
657 }
658 let av = (!a.is_null(i)).then(|| a.value(i));
659 let bv = (!b.is_null(i)).then(|| b.value(i));
660 if av != bv {
661 *is_changed = true;
662 }
663 }
664 }
665 Ok(changed.into_iter().map(|c| !c).collect())
666}
667
668fn apply_noise_with_retry(
691 orig: &RecordBatch,
692 noise_type: &str,
693 plan_cols: &[String],
694 rng: &mut Rng,
695 exclude_cols: &[String],
696) -> Result<(RecordBatch, Vec<bool>, Vec<String>), String> {
697 const MAX_NOISE_ATTEMPTS: usize = 4;
698 let (mut noisy, target_cols) =
699 apply_noise_to_batch(orig, noise_type, plan_cols, rng, exclude_cols)?;
700 if target_cols.is_empty() {
701 return Ok((noisy, vec![true; orig.num_rows()], target_cols));
702 }
703 let schema = noisy.schema();
704 let target_idxs: Vec<usize> = target_cols
705 .iter()
706 .filter_map(|c| schema.index_of(c).ok())
707 .collect();
708 let mut unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
709 let mut attempt = 1;
710 while unchanged.iter().any(|&u| u) && attempt < MAX_NOISE_ATTEMPTS {
711 let mut retry_rng = Rng::new(rng.next_u64());
712 let (retried, _) =
713 apply_noise_to_batch(orig, noise_type, plan_cols, &mut retry_rng, exclude_cols)?;
714 let mask = arrow::array::BooleanArray::from(unchanged.clone());
715 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
716 for &i in &target_idxs {
717 merged_cols[i] =
718 arrow::compute::kernels::zip::zip(&mask, retried.column(i), noisy.column(i))
719 .map_err(|e| format!("zip noise retry col {i}: {e}"))?;
720 }
721 noisy = RecordBatch::try_new(noisy.schema(), merged_cols)
722 .map_err(|e| format!("rebuild retried dup batch: {e}"))?;
723 unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
724 attempt += 1;
725 }
726 Ok((noisy, unchanged, target_cols))
727}
728
729fn apply_extra_pass_per_row(
756 noisy: &RecordBatch,
757 unchanged: &mut [bool],
758 all_types: &[&str],
759 rng: &mut Rng,
760 exclude_cols: &[String],
761) -> Result<RecordBatch, String> {
762 let n = noisy.num_rows();
763 if all_types.is_empty() || n == 0 {
764 return Ok(noisy.clone());
765 }
766 let per_row_type: Vec<usize> = (0..n).map(|_| rng.next_usize(all_types.len())).collect();
767
768 let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
779 for (i, &t) in per_row_type.iter().enumerate() {
780 groups[t].push(i as u32);
781 }
782
783 let schema = noisy.schema();
784 let mut pass_unchanged = vec![true; n];
785 let mut touched: HashMap<String, Vec<Option<String>>> = HashMap::new();
792
793 for (type_idx, idxs) in groups.into_iter().enumerate() {
794 if idxs.is_empty() {
795 continue;
796 }
797 let ttype = all_types[type_idx];
798 let idx_arr = UInt64Array::from_iter_values(idxs.iter().map(|&i| i as u64));
799 let subset = pick_rows(noisy, &idx_arr)?;
800 let (noised_subset, subset_unchanged, target_cols) =
801 apply_noise_with_retry(&subset, ttype, &[], rng, exclude_cols)?;
802 for col_name in &target_cols {
803 let Ok(col_idx) = schema.index_of(col_name) else {
804 continue;
805 };
806 let noised_col = noised_subset.column(col_idx).as_string::<i32>();
807 let buf = touched.entry(col_name.clone()).or_insert_with(|| {
808 let orig = noisy.column(col_idx).as_string::<i32>();
809 (0..n)
810 .map(|i| (!orig.is_null(i)).then(|| orig.value(i).to_string()))
811 .collect()
812 });
813 for (local_pos, &orig_i) in idxs.iter().enumerate() {
814 buf[orig_i as usize] = (!noised_col.is_null(local_pos))
815 .then(|| noised_col.value(local_pos).to_string());
816 }
817 }
818 for (local_pos, &orig_i) in idxs.iter().enumerate() {
819 pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
820 }
821 }
822
823 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
824 for (col_name, buf) in touched {
825 if let Ok(col_idx) = schema.index_of(&col_name) {
826 merged_cols[col_idx] = Arc::new(StringArray::from(buf));
827 }
828 }
829 let current = RecordBatch::try_new(noisy.schema(), merged_cols)
830 .map_err(|e| format!("rebuild extra-pass batch: {e}"))?;
831
832 for (u, pu) in unchanged.iter_mut().zip(pass_unchanged.iter()) {
833 *u = *u && *pu;
834 }
835 Ok(current)
836}
837
838fn topological_sort(
842 plans: &[EntityPlan],
843 name_to_idx: &HashMap<&str, usize>,
844) -> Result<Vec<usize>, String> {
845 let n = plans.len();
846 let mut in_degree = vec![0usize; n];
847 let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
848 for i in 0..n {
849 for remap in &plans[i].fk_remaps {
850 let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
851 continue;
852 };
853 if target == i {
854 continue;
855 }
856 adj[target].push(i); in_degree[i] += 1;
858 }
859 }
860 let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
861 let mut order = Vec::with_capacity(n);
862 let mut idx = 0;
863 while idx < queue.len() {
864 let v = queue[idx];
865 idx += 1;
866 order.push(v);
867 for &next in &adj[v] {
868 in_degree[next] -= 1;
869 if in_degree[next] == 0 {
870 queue.push(next);
871 }
872 }
873 }
874 if order.len() != n {
875 return Err("FK dependency cycle detected among entities".into());
876 }
877 Ok(order)
878}
879
880pub(crate) enum DatasetWriter {
884 Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
885 Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
886}
887
888impl DatasetWriter {
889 fn new(
890 output_format: &str,
891 path: &str,
892 schema: &Arc<Schema>,
893 metadata: &HashMap<String, String>,
894 ) -> Result<Self, String> {
895 let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
896 if output_format == "parquet" {
897 use parquet::basic::{Compression, ZstdLevel};
898 use parquet::file::properties::WriterProperties;
899 let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
900 let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
901 .iter()
902 .map(|(k, v)| parquet::file::metadata::KeyValue {
903 key: k.clone(),
904 value: Some(v.clone()),
905 })
906 .collect();
907 let props = WriterProperties::builder()
908 .set_compression(Compression::ZSTD(zstd))
909 .set_max_row_group_row_count(Some(1_000_000))
913 .set_key_value_metadata(Some(meta_kv))
914 .build();
915 let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
916 .map_err(|e| format!("ArrowWriter error: {e}"))?;
917 Ok(DatasetWriter::Parquet(Box::new(writer)))
918 } else {
919 let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
920 .map_err(|e| format!("IPC FileWriter error: {e}"))?;
921 Ok(DatasetWriter::Ipc(Box::new(writer)))
922 }
923 }
924
925 pub(crate) fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
926 match self {
927 DatasetWriter::Ipc(w) => w.write(batch).map_err(|e| format!("write ipc: {e}")),
928 DatasetWriter::Parquet(w) => w.write(batch).map_err(|e| format!("write parquet: {e}")),
929 }
930 }
931
932 fn finish(self) -> Result<(), String> {
933 match self {
934 DatasetWriter::Ipc(mut w) => w.finish().map_err(|e| format!("finish ipc: {e}")),
935 DatasetWriter::Parquet(w) => w
936 .close()
937 .map(|_| ())
938 .map_err(|e| format!("close parquet: {e}")),
939 }
940 }
941}
942
943pub fn run_pipeline(
946 ctx: &Context,
947 config: &PipelineConfig,
948 output_dir: &str,
949) -> Result<PipelineOutput, String> {
950 run_pipeline_with_progress(ctx, config, output_dir, None)
951}
952
953pub fn run_pipeline_with_progress(
960 ctx: &Context,
961 config: &PipelineConfig,
962 output_dir: &str,
963 mut progress: Option<&mut dyn FnMut(usize, usize)>,
964) -> Result<PipelineOutput, String> {
965 std::fs::create_dir_all(output_dir)
966 .map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
967 let t_start = std::time::Instant::now();
968
969 let t_alloc = t_start.elapsed().as_secs_f64();
970
971 let t1e = std::time::Instant::now();
975
976 let metadata = build_metadata_map(config);
982 let full_arc = Arc::new(build_full_schema(config, &metadata));
984 let dataset_ext = if config.output_format == "parquet" {
985 "parquet"
986 } else {
987 "ipc"
988 };
989 let dataset_path = format!("{}/{}.{}", output_dir, config.run_id, dataset_ext);
990 let mut writer =
991 DatasetWriter::new(&config.output_format, &dataset_path, &full_arc, &metadata)?;
992
993 let name_to_idx: HashMap<&str, usize> = config
995 .entity_plans
996 .iter()
997 .enumerate()
998 .map(|(i, p)| (p.name.as_str(), i))
999 .collect();
1000 let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;
1001
1002 let mut fk_pools: FkPoolMap = HashMap::new();
1004 let mut hn_pools: HashMap<String, HnPool> = HashMap::new();
1005
1006 let gt_ext = if config.output_format == "parquet" {
1010 "parquet"
1011 } else {
1012 "ipc"
1013 };
1014 let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
1015 let gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
1016 let mut gt_acc = crate::gt::GtAccumulator::new(>_draft_path)?;
1017
1018 let mut global_rid_offset = 0usize;
1019
1020 const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
1022 let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
1023 let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
1024 let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
1025 let mut graph_nodes_path: Option<String> = None;
1026 let mut graph_edges_path: Option<String> = None;
1027 if config.graph_enabled {
1028 let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
1029 let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
1030 node_writer = Some(
1031 crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
1032 .map_err(|e| format!("init node writer: {e}"))?,
1033 );
1034 edge_writer = Some(
1035 crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata)
1036 .map_err(|e| format!("init edge writer: {e}"))?,
1037 );
1038 graph_nodes_path = Some(nodes_ipc);
1039 graph_edges_path = Some(edges_ipc);
1040 }
1041
1042 let mut _t_write = 0.0f64;
1043 let mut _t_meta = 0.0f64;
1044 let mut _t_dup = 0.0f64;
1045 let mut _write_calls = 0u64;
1046 let hn_pool_sizing: std::collections::HashMap<&str, usize> =
1048 config
1049 .hard_neg_types
1050 .iter()
1051 .fold(std::collections::HashMap::new(), |mut m, hn| {
1052 let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
1053 *entry = (*entry).max(hn.count);
1054 m
1055 });
1056 let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
1057 let pool_only_names: std::collections::HashSet<&str> = config
1073 .entity_plans
1074 .iter()
1075 .filter(|p| p.pool_only)
1076 .map(|p| p.name.as_str())
1077 .collect();
1078
1079 let fk_targets: std::collections::HashSet<&str> = config
1080 .entity_plans
1081 .iter()
1082 .flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
1083 .collect();
1084
1085 for &plan_idx in &entity_order {
1086 let plan = &config.entity_plans[plan_idx];
1087 if plan.n_base == 0 {
1088 continue;
1089 }
1090
1091 let prefix = entity_prefix(plan_idx);
1092 let col_json_str = &plan.columns_json;
1093
1094 let fk_targeted = fk_targets.contains(plan.name.as_str());
1096 let mut fk_builder = arrow::array::StringBuilder::new();
1097 let mut fk_rid_builder = arrow::array::StringBuilder::new();
1098 let mut fk_count: usize = 0;
1099 let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
1104 let mut hn_slices: Vec<RecordBatch> = Vec::new();
1105 let mut hn_slice_rids: Vec<Vec<String>> = Vec::new();
1106 let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
1107 let mut fk_rng = Rng::new(config.seed.wrapping_add(42));
1108
1109 let batch_bounds: Vec<(usize, usize)> = (0..plan.n_base)
1115 .step_by(crate::entity_gen::BATCH_SIZE)
1116 .map(|offset| {
1117 (
1118 offset,
1119 (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE),
1120 )
1121 })
1122 .collect();
1123 let batch_weights: Vec<usize> = batch_bounds.iter().map(|&(_, bn)| bn).collect();
1124 let per_batch_noise_counts: Vec<Vec<usize>> = plan
1125 .noise_types
1126 .iter()
1127 .map(|n| distribute_by_weight(n.count, &batch_weights))
1128 .collect();
1129
1130 let mut col_lookup: Option<Vec<Option<usize>>> = None;
1132 let mut null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1134 let mut const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1135
1136 for (batch_idx, offset) in (0..plan.n_base)
1138 .step_by(crate::entity_gen::BATCH_SIZE)
1139 .enumerate()
1140 {
1141 let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
1142 let batch_seed = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);
1143
1144 let request_json = format!(
1146 r#"{{"entity_name":"{}","n":{},"seed":{},"columns":{},"row_offset":{}}}"#,
1147 plan.name, batch_n, batch_seed, col_json_str, offset,
1148 );
1149 let rb = crate::entity_gen::generate_entity_batch(ctx, &request_json)?;
1150
1151 let mut fk_edges_by_remap: Vec<(String, ArrayRef, f64)> = Vec::new();
1153 let rb = if !plan.fk_remaps.is_empty() {
1154 let mut r = rb;
1155 for remap in &plan.fk_remaps {
1156 if let Some(pool) = fk_pools.get(&remap.target_entity) {
1157 let want_edge = config.graph_enabled
1161 && !pool_only_names.contains(remap.target_entity.as_str());
1162 let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
1163 &r,
1164 pool,
1165 &remap.source_col,
1166 &mut fk_rng,
1167 want_edge,
1168 )?;
1169 r = remapped;
1170 if want_edge && let Some(rids) = target_rids {
1171 fk_edges_by_remap.push((
1172 remap.source_col.clone(),
1173 rids,
1174 remap.fk_weight,
1175 ));
1176 }
1177 }
1178 }
1179 r
1180 } else {
1181 rb
1182 };
1183
1184 if fk_targeted
1187 && fk_count < FK_POOL_CAP
1188 && let Some(ref id_col) = plan.identifier_col
1189 && let Some(col) = rb.column_by_name(id_col)
1190 {
1191 let s = col.as_string::<i32>();
1192 for i in 0..s.len() {
1193 if fk_count >= FK_POOL_CAP {
1194 break;
1195 }
1196 if !s.is_null(i) {
1197 fk_builder.append_value(s.value(i));
1198 if config.graph_enabled {
1199 fk_rid_builder.append_value(record_id_string(global_rid_offset + i));
1200 }
1201 fk_count += 1;
1202 }
1203 }
1204 }
1205
1206 if plan.pool_only {
1212 continue;
1213 }
1214
1215 let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1217 if hn_targeted && hn_accum < hn_max_pool {
1218 let take = (hn_max_pool - hn_accum).min(batch_n);
1219 hn_slices.push(rb.slice(0, take));
1220 if config.graph_enabled {
1224 hn_slice_rids.push(record_id_strs(global_rid_offset..global_rid_offset + take));
1225 }
1226 }
1227
1228 if col_lookup.is_none() {
1230 let plan_schema = rb.schema();
1231 col_lookup = Some(
1232 full_arc
1233 .fields()
1234 .iter()
1235 .skip(4)
1236 .map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1237 .collect(),
1238 );
1239 }
1240
1241 let batch_mids: Vec<String> = (offset..offset + batch_n)
1243 .map(|i| format!("{}-{}", prefix, pad_string(i)))
1244 .collect();
1245 let batch_mids_slice = &batch_mids[..];
1246
1247 let rid_slice = record_id_strs(global_rid_offset..global_rid_offset + batch_n);
1249 let t_m0 = std::time::Instant::now();
1250 let base_rb = add_metadata_and_align(
1251 &rb,
1252 &config.domain,
1253 &plan.name,
1254 &rid_slice,
1255 batch_mids_slice,
1256 &full_arc,
1257 col_lookup.as_ref().unwrap(),
1258 &mut null_cache,
1259 &mut const_arr_cache,
1260 )?;
1261 _t_meta += t_m0.elapsed().as_secs_f64();
1262
1263 if config.graph_enabled
1265 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1266 {
1267 nw.write_batch(&base_rb)
1268 .map_err(|e| format!("write node: {e}"))?;
1269 let fk_edges_by_remap: Vec<_> = fk_edges_by_remap
1270 .iter()
1271 .map(|(subtype, target_rids, weight)| {
1272 (subtype, target_rids.as_string::<i32>(), *weight)
1273 })
1274 .collect();
1275 for (i, rid) in rid_slice.iter().enumerate().take(batch_n) {
1276 for (subtype, target_rids, weight) in &fk_edges_by_remap {
1277 ew.push(rid, target_rids.value(i), "fk", subtype, *weight)
1278 .map_err(|e| format!("push fk edge: {e}"))?;
1279 }
1280 }
1281 }
1282
1283 let t_w0 = std::time::Instant::now();
1284 writer
1285 .write(&base_rb)
1286 .map_err(|e| format!("write base: {e}"))?;
1287 _t_write += t_w0.elapsed().as_secs_f64();
1288 _write_calls += 1;
1289
1290 gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;
1292
1293 global_rid_offset += batch_n;
1294 if let Some(cb) = &mut progress {
1295 cb(global_rid_offset, config.size);
1296 }
1297
1298 let has_dups: bool = per_batch_noise_counts
1306 .iter()
1307 .any(|counts| counts[batch_idx] > 0);
1308 if has_dups {
1309 let t_d0 = std::time::Instant::now();
1310 let mut dup_batches: Vec<RecordBatch> = Vec::new();
1311 let mut dup_mids_buf: Vec<String> = Vec::new();
1312
1313 let fk_exclude_cols: Vec<String> = plan
1315 .fk_remaps
1316 .iter()
1317 .map(|r| r.source_col.clone())
1318 .collect();
1319
1320 let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
1322 .noise_types
1323 .iter()
1324 .enumerate()
1325 .filter_map(|(ni, n)| {
1326 let cnt = per_batch_noise_counts[ni][batch_idx];
1327 if cnt == 0 {
1328 return None;
1329 }
1330 let mut b = UInt64Builder::with_capacity(cnt);
1331 for _ in 0..cnt {
1332 b.append_value(batch_rng.next_usize(batch_n) as u64);
1333 }
1334 Some((
1335 b.finish(),
1336 batch_rng.next_u64(),
1337 n.noise_type.as_str(),
1338 n.columns.as_slice(),
1339 cnt,
1340 ))
1341 })
1342 .collect();
1343
1344 let all_types: Vec<&str> = plan
1348 .noise_types
1349 .iter()
1350 .map(|n| n.noise_type.as_str())
1351 .collect();
1352 let all_types_ref = &all_types;
1353 let noise_passes = config.noise_passes.max(1);
1354
1355 type DupNoiseResult = Result<(RecordBatch, Vec<String>, Vec<bool>), String>;
1356 let fk_exclude_ref = &fk_exclude_cols;
1357 let rb_ref = &rb;
1358 let prefix_ref: &str = &prefix;
1359 let results: Vec<DupNoiseResult> = ndata
1372 .par_iter()
1373 .map(|(indices, seed, ntype, cols, cnt)| {
1374 let dup = pick_rows(rb_ref, indices)?;
1375 let mut rng = Rng::new(*seed);
1376
1377 let (mut noisy, mut unchanged, _) =
1383 apply_noise_with_retry(&dup, ntype, cols, &mut rng, fk_exclude_ref)?;
1384
1385 for _ in 1..noise_passes {
1397 noisy = apply_extra_pass_per_row(
1398 &noisy,
1399 &mut unchanged,
1400 all_types_ref,
1401 &mut rng,
1402 fk_exclude_ref,
1403 )?;
1404 }
1405
1406 let mut mb = Vec::with_capacity(*cnt);
1407 for j in 0..*cnt {
1408 let global_idx = offset + indices.value(j) as usize;
1413 mb.push(format!("{}-{}", prefix_ref, pad_string(global_idx)));
1414 }
1415 Ok((noisy, mb, unchanged))
1416 })
1417 .collect();
1418 let mut dup_is_identical_buf: Vec<bool> = Vec::new();
1419 for res in results {
1420 let (rb, mb, unchanged) = res?;
1421 dup_batches.push(rb);
1422 dup_mids_buf.extend(mb);
1423 dup_is_identical_buf.extend(unchanged);
1424 }
1425 _t_dup += t_d0.elapsed().as_secs_f64();
1426
1427 if !dup_batches.is_empty() {
1429 let dup_total = dup_mids_buf.len();
1430 let dup_rids = record_id_strs(global_rid_offset..global_rid_offset + dup_total);
1431
1432 let concated = if dup_batches.len() == 1 {
1433 dup_batches.into_iter().next().unwrap()
1434 } else {
1435 let schema = dup_batches[0].schema();
1436 let n_fields = dup_batches[0].num_columns();
1437 let mut concat_arrays = Vec::with_capacity(n_fields);
1438 for i in 0..n_fields {
1439 let refs: Vec<&dyn Array> =
1440 dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
1441 concat_arrays.push(
1442 arrow::compute::concat(&refs)
1443 .map_err(|e| format!("concat dup col {i}: {e}"))?,
1444 );
1445 }
1446 RecordBatch::try_new(schema, concat_arrays)
1447 .map_err(|e| format!("concat dups: {e}"))?
1448 };
1449
1450 let dup_rb_full = add_metadata_and_align(
1451 &concated,
1452 &config.domain,
1453 &plan.name,
1454 &dup_rids,
1455 &dup_mids_buf,
1456 &full_arc,
1457 col_lookup.as_ref().unwrap(),
1458 &mut null_cache,
1459 &mut const_arr_cache,
1460 )?;
1461
1462 let t_wd = std::time::Instant::now();
1463 writer
1464 .write(&dup_rb_full)
1465 .map_err(|e| format!("write dups: {e}"))?;
1466 _t_write += t_wd.elapsed().as_secs_f64();
1467 _write_calls += 1;
1468
1469 if config.graph_enabled
1471 && let Some(nw) = node_writer.as_mut()
1472 {
1473 nw.write_batch(&dup_rb_full)
1474 .map_err(|e| format!("write dup node: {e}"))?;
1475 }
1476
1477 let is_identical_arr: ArrayRef =
1478 Arc::new(arrow::array::BooleanArray::from(dup_is_identical_buf));
1479 gt_acc.push_dup_batch(
1480 dup_rb_full.column(0),
1481 dup_rb_full.column(2),
1482 dup_rb_full.column(3),
1483 &is_identical_arr,
1484 )?;
1485 global_rid_offset += dup_total;
1486 }
1487 }
1488 }
1489
1490 if let Some(id_col) = &plan.identifier_col
1492 && fk_builder.len() > 0
1493 {
1494 let arr = Arc::new(fk_builder.finish()) as ArrayRef;
1495 if config.graph_enabled {
1496 let rid_arr = Arc::new(fk_rid_builder.finish()) as ArrayRef;
1500 let schema = Arc::new(Schema::new(vec![
1501 Field::new(id_col, DataType::Utf8, true),
1502 Field::new("record_id", DataType::Utf8, true),
1503 ]));
1504 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
1505 fk_pools.insert(plan.name.clone(), pool_rb);
1506 }
1507 } else {
1508 let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
1509 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
1510 fk_pools.insert(plan.name.clone(), pool_rb);
1511 }
1512 }
1513 }
1514
1515 if !hn_slices.is_empty() {
1517 let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1518 let batch = if hn_slices.len() == 1 {
1519 let only = hn_slices.into_iter().next().unwrap();
1526 let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
1527 pick_rows(&only, &idx)?
1528 } else {
1529 let schema = hn_slices[0].schema();
1530 let n_fields = schema.fields().len();
1531 let mut concat_arrays = Vec::with_capacity(n_fields);
1532 for i in 0..n_fields {
1533 let refs: Vec<&dyn Array> =
1534 hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
1535 concat_arrays.push(
1536 arrow::compute::concat(&refs)
1537 .map_err(|e| format!("hn pool concat col {i}: {e}"))?,
1538 );
1539 }
1540 RecordBatch::try_new(schema, concat_arrays)
1541 .map_err(|e| format!("hn pool concat batch: {e}"))?
1542 };
1543 let record_ids = if config.graph_enabled {
1545 hn_slice_rids.into_iter().flatten().collect::<Vec<String>>()
1546 } else {
1547 Vec::new()
1548 };
1549 hn_pools.insert(
1550 plan.name.clone(),
1551 HnPool {
1552 batch,
1553 total_count,
1554 record_ids,
1555 },
1556 );
1557 }
1558 }
1559 let t1e_elapsed = t1e.elapsed().as_secs_f64();
1560 log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
1561
1562 let t2 = std::time::Instant::now();
1564 let mut hn_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1566 let mut hn_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1567 let mut hn_master_id_counter: u64 = 0;
1575 for hn_cfg in config.hard_neg_types.iter() {
1576 if hn_cfg.count == 0 {
1577 continue;
1578 }
1579 let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
1580 continue;
1581 };
1582 let hn_total = pool_data.total_count;
1583
1584 let pool_rb = &pool_data.batch;
1585 let n_hn_max = hn_cfg.count.min(hn_total / 2);
1586 if n_hn_max == 0 {
1587 continue;
1588 }
1589
1590 let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
1591 pool_rb,
1592 &hn_cfg.config_json,
1593 n_hn_max,
1594 config.seed.wrapping_add(200),
1595 config.graph_enabled,
1596 )?;
1597
1598 let n_hn = hn_rb.num_rows();
1599 if n_hn == 0 {
1600 continue;
1601 }
1602
1603 let hn_rids = record_id_strs(global_rid_offset..global_rid_offset + n_hn);
1604
1605 let hn_mids: Vec<String> = (0..n_hn)
1606 .map(|_i| {
1607 let id = hn_master_id_counter;
1608 hn_master_id_counter += 1;
1609 format!("HN-{:09}", id)
1610 })
1611 .collect();
1612
1613 let hn_schema = hn_rb.schema();
1615 let hn_col_lookup: Vec<Option<usize>> = full_arc
1616 .fields()
1617 .iter()
1618 .skip(4)
1619 .map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1620 .collect();
1621
1622 let hn_rb_full = add_metadata_and_align(
1623 &hn_rb,
1624 &config.domain,
1625 &hn_cfg.entity_type,
1626 &hn_rids,
1627 &hn_mids,
1628 &full_arc,
1629 &hn_col_lookup,
1630 &mut hn_null_cache,
1631 &mut hn_const_arr_cache,
1632 )?;
1633
1634 let t_wh = std::time::Instant::now();
1635 writer
1636 .write(&hn_rb_full)
1637 .map_err(|e| format!("write hn: {e}"))?;
1638 _t_write += t_wh.elapsed().as_secs_f64();
1639 _write_calls += 1;
1640
1641 if config.graph_enabled
1644 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1645 {
1646 nw.write_batch(&hn_rb_full)
1647 .map_err(|e| format!("write hn node: {e}"))?;
1648 if let Some((idx_a, pattern)) = hn_src {
1649 for i in 0..n_hn {
1650 let tgt = &pool_data.record_ids[idx_a[i]];
1651 ew.push(&hn_rids[i], tgt, "hard_neg", &pattern, 1.0)
1652 .map_err(|e| format!("push hn edge: {e}"))?;
1653 }
1654 }
1655 }
1656
1657 gt_acc.push_other_batch(
1659 hn_rb_full.column(0),
1660 hn_rb_full.column(2),
1661 hn_rb_full.column(3),
1662 )?;
1663
1664 global_rid_offset += n_hn;
1665 }
1666 let t2_elapsed = t2.elapsed().as_secs_f64();
1667
1668 {
1670 let mut canary_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1671 let mut canary_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1672 crate::canary::generate_all(
1673 ctx,
1674 config,
1675 &full_arc,
1676 &mut canary_null_cache,
1677 &mut canary_const_arr_cache,
1678 &mut global_rid_offset,
1679 &fk_pools,
1680 &mut writer,
1681 &mut node_writer,
1682 &mut gt_acc,
1683 )?;
1684 }
1685
1686 let t3 = std::time::Instant::now();
1688 writer
1689 .finish()
1690 .map_err(|e| format!("finish dataset writer: {e}"))?;
1691 let t3a_elapsed = t3.elapsed().as_secs_f64();
1692
1693 let t3b = std::time::Instant::now();
1694 let t_gt0 = std::time::Instant::now();
1695 let crate::gt::GtResult {
1696 n_exact_dup,
1697 n_fuzzy_dup,
1698 n_hard_neg,
1699 n_unique,
1700 n_masters,
1701 cluster_map,
1702 } = gt_acc.finish(
1703 config.difficulty.as_str(),
1704 &config.output_format,
1705 >_path,
1706 &metadata,
1707 config.graph_enabled,
1708 )?;
1709 let _t_gt_compute = 0.0f64;
1710 let _t_gt_write = t_gt0.elapsed().as_secs_f64();
1711 let t3b_elapsed = t3b.elapsed().as_secs_f64();
1712 let _t3_elapsed = t3.elapsed().as_secs_f64();
1713
1714 if config.graph_enabled
1716 && let Some(ew) = edge_writer.as_mut()
1717 {
1718 crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
1719 .map_err(|e| format!("push dup clusters: {e}"))?;
1720 }
1721
1722 let duration = t_start.elapsed().as_secs_f64();
1723
1724 log::debug!(
1725 "[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",
1726 t_alloc,
1727 t1e_elapsed,
1728 t2_elapsed,
1729 t3a_elapsed,
1730 _t_meta,
1731 _t_dup,
1732 _t_write,
1733 _t_gt_compute,
1734 _t_gt_write,
1735 t3b_elapsed,
1736 duration,
1737 );
1738
1739 let stats = PipelineStats {
1740 total_records: global_rid_offset,
1741 exact_dups: n_exact_dup,
1742 fuzzy_dups: n_fuzzy_dup,
1743 hard_negs: n_hard_neg,
1744 uniques: n_unique,
1745 masters: n_masters,
1746 };
1747
1748 let mut graph_nodes_final: Option<String> = None;
1750 let mut graph_edges_final: Option<String> = None;
1751 if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
1752 nw.finish()
1753 .map_err(|e| format!("finish node writer: {e}"))?;
1754 ew.finish()
1755 .map_err(|e| format!("finish edge writer: {e}"))?;
1756 let mut np = graph_nodes_path.clone().unwrap();
1757 let mut ep = graph_edges_path.clone().unwrap();
1758 if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
1759 let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
1760 let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
1761 convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
1762 convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
1763 std::fs::remove_file(&np).ok();
1764 std::fs::remove_file(&ep).ok();
1765 np = np_pq;
1766 ep = ep_pq;
1767 }
1768 graph_nodes_final = Some(np);
1769 graph_edges_final = Some(ep);
1770 }
1771
1772 Ok(PipelineOutput {
1773 output_files: vec![dataset_path],
1774 gt_file: gt_path,
1775 stats,
1776 nodes: graph_nodes_final,
1777 edges: graph_edges_final,
1778 })
1779}
1780
1781fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
1785 use arrow::ipc::reader::FileReader;
1786
1787 let ipc_file =
1788 std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
1789 let ipc_reader =
1790 FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
1791 let schema = ipc_reader.schema();
1792
1793 let parquet_file =
1794 std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
1795 let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
1796 let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
1797 .metadata()
1798 .iter()
1799 .map(|(k, v)| parquet::file::metadata::KeyValue {
1800 key: k.clone(),
1801 value: Some(v.clone()),
1802 })
1803 .collect();
1804 let props = parquet::file::properties::WriterProperties::builder()
1805 .set_compression(parquet::basic::Compression::ZSTD(zstd_level))
1806 .set_max_row_group_row_count(Some(1_000_000))
1807 .set_key_value_metadata(Some(meta_kv))
1808 .build();
1809 let mut parquet_writer =
1810 parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
1811 .map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;
1812
1813 for batch_result in ipc_reader {
1814 let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
1815 parquet_writer
1816 .write(&batch)
1817 .map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
1818 }
1819 parquet_writer
1820 .close()
1821 .map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
1822 Ok(())
1823}
1824
1825fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
1828 let ts = std::time::SystemTime::now()
1829 .duration_since(std::time::UNIX_EPOCH)
1830 .map(|d| d.as_secs().to_string())
1831 .unwrap_or_else(|_| "unknown".into());
1832 HashMap::from([
1833 (
1834 "dupehell.generator".into(),
1835 format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
1836 ),
1837 (
1838 "dupehell.provenance".into(),
1839 "dupehell-synthetic-data".into(),
1840 ),
1841 ("dupehell.license".into(), "MIT".into()),
1842 (
1843 "dupehell.purpose".into(),
1844 "Educational Use Only — Record Linkage Benchmarking".into(),
1845 ),
1846 (
1847 "dupehell.url".into(),
1848 "https://github.com/vntoinekaio/DupeHell".into(),
1849 ),
1850 ("dupehell.domain".into(), config.domain.clone()),
1851 ("dupehell.size".into(), config.size.to_string()),
1852 ("dupehell.seed".into(), config.seed.to_string()),
1853 ("dupehell.difficulty".into(), config.difficulty.clone()),
1854 ("dupehell.run_id".into(), config.run_id.clone()),
1855 ("dupehell.timestamp".into(), ts),
1856 (
1857 "dupehell.structural_key_columns".into(),
1858 structural_key_columns_note(config),
1859 ),
1860 ])
1861}
1862
1863fn structural_key_columns_note(config: &PipelineConfig) -> String {
1868 let cols: Vec<&str> = config
1869 .entity_plans
1870 .iter()
1871 .filter_map(|p| p.identifier_col.as_deref())
1872 .collect();
1873 format!(
1874 "{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
1875 cols.join(", ")
1876 )
1877}
1878
1879fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
1883 let mut field_map: Vec<(String, DataType, bool)> = Vec::new();
1884 let metadata_fields = ["record_id", "domain", "entity_type", "master_id"];
1885 for mf in &metadata_fields {
1886 field_map.push((mf.to_string(), DataType::Utf8, false));
1887 }
1888 for plan in &config.entity_plans {
1889 if plan.pool_only {
1894 continue;
1895 }
1896 let cols: Vec<serde_json::Value> =
1897 serde_json::from_str(&plan.columns_json).unwrap_or_default();
1898 for col in &cols {
1899 let name = col["name"].as_str().unwrap_or("").to_string();
1900 if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
1901 continue;
1902 }
1903 let col_type = col_type_from_request(col);
1904 field_map.push((name, col_type, true));
1905 }
1906 }
1907 let fields: Vec<Field> = field_map
1908 .into_iter()
1909 .map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
1910 .collect();
1911 Schema::new(fields).with_metadata(metadata.clone())
1912}
1913
1914fn col_type_from_request(v: &serde_json::Value) -> DataType {
1916 match v["type"].as_str() {
1917 Some("int") => DataType::Int64,
1918 Some("float") => DataType::Float64,
1919 Some("boolean") => DataType::Boolean,
1920 Some("date") | Some("datetime") => DataType::Utf8,
1921 _ => DataType::Utf8,
1922 }
1923}
1924
1925#[allow(clippy::too_many_arguments)]
1928pub(crate) fn add_metadata_and_align(
1929 rb: &RecordBatch,
1930 domain: &str,
1931 entity_type: &str,
1932 record_ids: &[String],
1933 master_ids: &[String],
1934 full_arc: &Arc<Schema>,
1935 col_lookup: &[Option<usize>],
1936 null_cache: &mut FxHashMap<(DataType, usize), ArrayRef>,
1937 const_arr_cache: &mut FxHashMap<(String, usize), ArrayRef>,
1938) -> Result<RecordBatch, String> {
1939 let n = rb.num_rows();
1940 let domain_arr = const_arr_cache
1946 .entry((domain.to_string(), n))
1947 .or_insert_with(|| {
1948 Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
1949 domain, n,
1950 ))) as ArrayRef
1951 })
1952 .clone();
1953 let et_arr = const_arr_cache
1954 .entry((entity_type.to_string(), n))
1955 .or_insert_with(|| {
1956 Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
1957 entity_type,
1958 n,
1959 ))) as ArrayRef
1960 })
1961 .clone();
1962 let rid_arr = Arc::new(StringArray::from_iter_values(
1963 record_ids.iter().map(|s| s.as_str()),
1964 )) as ArrayRef;
1965 let mid_arr = Arc::new(StringArray::from_iter_values(
1966 master_ids.iter().map(|s| s.as_str()),
1967 )) as ArrayRef;
1968
1969 let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
1970 all_arrays.push(rid_arr);
1971 all_arrays.push(domain_arr);
1972 all_arrays.push(et_arr);
1973 all_arrays.push(mid_arr);
1974
1975 for (i, maybe_idx) in col_lookup.iter().enumerate() {
1978 match maybe_idx {
1979 Some(idx) => all_arrays.push(rb.column(*idx).clone()),
1980 None => {
1981 let dt = full_arc.field(i + 4).data_type();
1982 let arr = null_cache
1983 .entry((dt.clone(), n))
1984 .or_insert_with(|| arrow::array::new_null_array(dt, n));
1985 all_arrays.push(Arc::clone(arr));
1986 }
1987 }
1988 }
1989
1990 RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
1992 .map_err(|e| format!("add_metadata_and_align RecordBatch: {e}"))
1993}
1994
1995fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
1996 use arrow::compute::take;
1997 let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
1998 .map(|i| take(rb.column(i), indices, None))
1999 .collect::<Result<Vec<_>, _>>()
2000 .map_err(|e| format!("take error: {e}"))?;
2001
2002 RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
2003}
2004
2005#[cfg(test)]
2006mod tests {
2007 use super::*;
2008 use crate::schema::{build_pipeline_config, load_schema};
2009
2010 #[test]
2017 fn test_no_person_company_name_cross_contamination() {
2018 for col in [
2019 "given_name",
2020 "family_name",
2021 "middle_name",
2022 "first_name",
2023 "last_name",
2024 ] {
2025 assert!(
2026 noise_type_targets_column("names", col),
2027 "{col} should still match 'names'"
2028 );
2029 assert!(
2030 !noise_type_targets_column("companies", col),
2031 "{col} should NOT match 'companies' (person name, not a company name)"
2032 );
2033 }
2034 for col in ["company_name", "legal_name", "trading_name"] {
2035 assert!(
2036 noise_type_targets_column("companies", col),
2037 "{col} should still match 'companies'"
2038 );
2039 assert!(
2040 !noise_type_targets_column("names", col),
2041 "{col} should NOT match 'names' (company name, not a person name)"
2042 );
2043 }
2044 assert!(
2045 !noise_type_targets_column("addresses", "email_address"),
2046 "email_address should NOT match 'addresses' (it's an email, not a postal address)"
2047 );
2048 assert!(
2049 noise_type_targets_column("identifiers", "email_address"),
2050 "email_address should still match 'identifiers' (corrupt_email)"
2051 );
2052 assert!(
2053 noise_type_targets_column("addresses", "residential_address_line1"),
2054 "a real postal address column should still match 'addresses'"
2055 );
2056 }
2057
2058 fn manifest_dir() -> std::path::PathBuf {
2059 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).to_path_buf()
2060 }
2061
2062 fn scratch_dir(name: &str) -> std::path::PathBuf {
2065 let dir = manifest_dir()
2066 .join("target")
2067 .join("pipeline_test_tmp")
2068 .join(format!("{name}_{}", std::process::id()));
2069 let _ = std::fs::remove_dir_all(&dir);
2070 std::fs::create_dir_all(&dir).expect("create scratch dir");
2071 dir
2072 }
2073
2074 fn run_kyc(
2075 output_dir: &std::path::Path,
2076 output_format: &str,
2077 graph_enabled: bool,
2078 ) -> PipelineOutput {
2079 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2080 let mut ctx = Context::new(
2081 "kyc",
2082 "en",
2083 &manifest_dir().join("assets/pools").to_string_lossy(),
2084 )
2085 .expect("load context");
2086 let config = build_pipeline_config(
2087 "kyc",
2088 200,
2089 42,
2090 "medium",
2091 0.1,
2092 0.3,
2093 &schema,
2094 "pipeline_test_run",
2095 output_format,
2096 graph_enabled,
2097 "parquet",
2098 None,
2099 )
2100 .expect("build config");
2101 ctx.enable_watermark(&config.domain, config.size, config.seed);
2102 run_pipeline(&ctx, &config, &output_dir.to_string_lossy()).expect("run_pipeline")
2103 }
2104
2105 #[test]
2106 fn test_run_pipeline_produces_dataset_and_gt() {
2107 let dir = scratch_dir("basic");
2108 let output = run_kyc(&dir, "parquet", false);
2109
2110 assert_eq!(output.output_files.len(), 1);
2111 for f in &output.output_files {
2112 assert!(std::path::Path::new(f).exists(), "missing dataset file {f}");
2113 }
2114 assert!(std::path::Path::new(&output.gt_file).exists());
2115 assert!(output.nodes.is_none());
2116 assert!(output.edges.is_none());
2117
2118 let s = &output.stats;
2123 assert!(s.total_records >= s.exact_dups + s.fuzzy_dups + s.hard_negs + s.uniques);
2124 assert!(s.total_records > 0);
2125 assert!(s.masters > 0);
2126
2127 let _ = std::fs::remove_dir_all(&dir);
2128 }
2129
2130 #[test]
2131 fn test_run_pipeline_ipc_format() {
2132 let dir = scratch_dir("ipc");
2133 let output = run_kyc(&dir, "ipc", false);
2134 assert!(output.output_files[0].ends_with(".ipc"));
2135 assert!(std::path::Path::new(&output.output_files[0]).exists());
2136 let _ = std::fs::remove_dir_all(&dir);
2137 }
2138
2139 #[test]
2140 fn test_run_pipeline_graph_output() {
2141 let dir = scratch_dir("graph");
2142 let output = run_kyc(&dir, "parquet", true);
2143 let nodes = output.nodes.expect("nodes path");
2144 let edges = output.edges.expect("edges path");
2145 assert!(std::path::Path::new(&nodes).exists());
2146 assert!(std::path::Path::new(&edges).exists());
2147 let _ = std::fs::remove_dir_all(&dir);
2148 }
2149
2150 #[test]
2151 fn test_fk_edge_weight_from_schema_and_default() {
2152 let dir = scratch_dir("fk_weight");
2159 let schema =
2160 load_schema("aviation", &manifest_dir().join("schemas")).expect("load aviation.json");
2161 let mut ctx = Context::new(
2162 "aviation",
2163 "en",
2164 &manifest_dir().join("assets/pools").to_string_lossy(),
2165 )
2166 .expect("load context");
2167 let config = build_pipeline_config(
2168 "aviation",
2169 300,
2170 42,
2171 "medium",
2172 0.1,
2173 0.3,
2174 &schema,
2175 "pipeline_test_fk_weight",
2176 "ipc",
2177 true,
2178 "ipc",
2179 None,
2180 )
2181 .expect("build config");
2182 ctx.enable_watermark(&config.domain, config.size, config.seed);
2183 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2184
2185 let edges_path = output.edges.expect("edges path");
2186 let file = std::fs::File::open(&edges_path).expect("open edges file");
2187 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2188
2189 let mut subtypes_seen: std::collections::HashMap<String, f64> =
2190 std::collections::HashMap::new();
2191 for batch in reader {
2192 let batch = batch.expect("read edges batch");
2193 let edge_type = batch
2194 .column_by_name("edge_type")
2195 .expect("edge_type col")
2196 .as_any()
2197 .downcast_ref::<arrow::array::StringArray>()
2198 .expect("edge_type as string");
2199 let subtype = batch
2200 .column_by_name("subtype")
2201 .expect("subtype col")
2202 .as_any()
2203 .downcast_ref::<arrow::array::StringArray>()
2204 .expect("subtype as string");
2205 let weight = batch
2206 .column_by_name("weight")
2207 .expect("weight col")
2208 .as_any()
2209 .downcast_ref::<arrow::array::Float64Array>()
2210 .expect("weight as f64");
2211 for i in 0..batch.num_rows() {
2212 if edge_type.value(i) == "fk" {
2213 subtypes_seen.insert(subtype.value(i).to_string(), weight.value(i));
2214 }
2215 }
2216 }
2217
2218 assert_eq!(
2219 subtypes_seen.get("aircraft_id"),
2220 Some(&0.5),
2221 "flight->aircraft fk_weight from schema not propagated to edge weight"
2222 );
2223 assert_eq!(
2224 subtypes_seen.get("airline_id"),
2225 Some(&1.0),
2226 "flight->airline fk_weight (schema default) not propagated to edge weight"
2227 );
2228
2229 let _ = std::fs::remove_dir_all(&dir);
2230 }
2231
2232 #[test]
2239 fn test_run_pipeline_dup_cluster_edges_well_formed() {
2240 let dir = scratch_dir("dup_cluster_edges");
2241 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2246 let mut ctx = Context::new(
2247 "kyc",
2248 "en",
2249 &manifest_dir().join("assets/pools").to_string_lossy(),
2250 )
2251 .expect("load context");
2252 let config = build_pipeline_config(
2253 "kyc",
2254 200,
2255 42,
2256 "medium",
2257 0.1,
2258 0.3,
2259 &schema,
2260 "pipeline_test_dup_cluster_edges",
2261 "ipc",
2262 true,
2263 "ipc",
2264 None,
2265 )
2266 .expect("build config");
2267 ctx.enable_watermark(&config.domain, config.size, config.seed);
2268 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2269
2270 let edges_path = output.edges.expect("edges path");
2271 let file = std::fs::File::open(&edges_path).expect("open edges file");
2272 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2273
2274 let mut n_exact = 0usize;
2275 let mut n_fuzzy = 0usize;
2276 for batch in reader {
2277 let batch = batch.expect("read edges batch");
2278 let src = batch
2279 .column_by_name("source_node_id")
2280 .expect("source_node_id col")
2281 .as_any()
2282 .downcast_ref::<arrow::array::StringArray>()
2283 .expect("source_node_id as string");
2284 let tgt = batch
2285 .column_by_name("target_node_id")
2286 .expect("target_node_id col")
2287 .as_any()
2288 .downcast_ref::<arrow::array::StringArray>()
2289 .expect("target_node_id as string");
2290 let edge_type = batch
2291 .column_by_name("edge_type")
2292 .expect("edge_type col")
2293 .as_any()
2294 .downcast_ref::<arrow::array::StringArray>()
2295 .expect("edge_type as string");
2296 let weight = batch
2297 .column_by_name("weight")
2298 .expect("weight col")
2299 .as_any()
2300 .downcast_ref::<arrow::array::Float64Array>()
2301 .expect("weight as f64");
2302
2303 for i in 0..batch.num_rows() {
2304 match edge_type.value(i) {
2305 "exact_dup" => n_exact += 1,
2306 "fuzzy_dup" => n_fuzzy += 1,
2307 _ => continue,
2308 }
2309 assert!(
2314 parse_record_idx(src.value(i)).is_some(),
2315 "malformed source_node_id {:?} on a dup-cluster edge",
2316 src.value(i)
2317 );
2318 assert!(
2319 parse_record_idx(tgt.value(i)).is_some(),
2320 "malformed target_node_id {:?} on a dup-cluster edge",
2321 tgt.value(i)
2322 );
2323 assert!((weight.value(i) - 1.0).abs() < 1e-9);
2324 }
2325 }
2326
2327 assert!(n_exact + n_fuzzy > 0, "no dup-cluster edges emitted");
2332 assert!(output.stats.exact_dups + output.stats.fuzzy_dups > 0);
2333
2334 let _ = std::fs::remove_dir_all(&dir);
2335 }
2336
2337 #[test]
2338 fn test_run_pipeline_deterministic_stats() {
2339 let dir_a = scratch_dir("det_a");
2340 let dir_b = scratch_dir("det_b");
2341 let a = run_kyc(&dir_a, "parquet", false);
2342 let b = run_kyc(&dir_b, "parquet", false);
2343 assert_eq!(a.stats.total_records, b.stats.total_records);
2344 assert_eq!(a.stats.exact_dups, b.stats.exact_dups);
2345 assert_eq!(a.stats.fuzzy_dups, b.stats.fuzzy_dups);
2346 assert_eq!(a.stats.hard_negs, b.stats.hard_negs);
2347 assert_eq!(a.stats.uniques, b.stats.uniques);
2348 assert_eq!(a.stats.masters, b.stats.masters);
2349 let _ = std::fs::remove_dir_all(&dir_a);
2350 let _ = std::fs::remove_dir_all(&dir_b);
2351 }
2352
2353 #[test]
2354 fn test_distribute_by_weight_sums_to_total_and_is_proportional() {
2355 let weights = [500_000usize, 216_666];
2363 let counts = distribute_by_weight(283_334, &weights);
2364 assert_eq!(counts.len(), 2);
2365 assert_eq!(counts.iter().sum::<usize>(), 283_334);
2366 assert!(counts[0] > 0, "first batch got zero duplicates");
2370 assert!(counts[1] > 0, "last batch got zero duplicates");
2371 assert!(counts[0] > counts[1]);
2372
2373 assert_eq!(distribute_by_weight(0, &weights), vec![0, 0]);
2375 assert_eq!(distribute_by_weight(100, &[0, 0]), vec![0, 0]);
2376 }
2377
2378 #[test]
2379 fn test_entity_prefix_unique_per_index() {
2380 let a = entity_prefix(0);
2381 let b = entity_prefix(1);
2382 assert_ne!(a, b);
2383 assert_eq!(a, "E00000");
2384 assert_eq!(entity_prefix(42), "E00042");
2385 }
2386
2387 #[test]
2388 fn test_run_pipeline_hell_multi_pass_touches_multiple_columns() {
2389 let dir = scratch_dir("hell_multipass");
2398 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2399 let mut ctx = Context::new(
2400 "kyc",
2401 "en",
2402 &manifest_dir().join("assets/pools").to_string_lossy(),
2403 )
2404 .expect("load context");
2405 let config = build_pipeline_config(
2406 "kyc",
2407 5000,
2408 7,
2409 "hell",
2410 0.1,
2411 0.3,
2412 &schema,
2413 "pipeline_test_hell_multipass",
2414 "ipc",
2415 false,
2416 "ipc",
2417 None,
2418 )
2419 .expect("build config");
2420 ctx.enable_watermark(&config.domain, config.size, config.seed);
2421 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2422
2423 let read_all = |path: &str| -> RecordBatch {
2424 let file = std::fs::File::open(path).unwrap();
2425 let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
2426 let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2427 let schema = batches[0].schema();
2428 let refs: Vec<Vec<&dyn Array>> = (0..schema.fields().len())
2429 .map(|i| batches.iter().map(|b| b.column(i).as_ref()).collect())
2430 .collect();
2431 let cols: Vec<ArrayRef> = refs
2432 .iter()
2433 .map(|r| arrow::compute::concat(r).unwrap())
2434 .collect();
2435 RecordBatch::try_new(schema, cols).unwrap()
2436 };
2437 let ds = read_all(&output.output_files[0]);
2438 let gt = read_all(&output.gt_file);
2439
2440 let ds_rid = ds.column_by_name("record_id").unwrap().as_string::<i32>();
2441 let ds_mid = ds.column_by_name("master_id").unwrap().as_string::<i32>();
2442 let gt_rid = gt.column_by_name("record_id").unwrap().as_string::<i32>();
2443 let gt_mt = gt.column_by_name("match_type").unwrap().as_string::<i32>();
2444
2445 let mut match_type_by_rid: HashMap<&str, &str> = HashMap::new();
2446 for i in 0..gt.num_rows() {
2447 match_type_by_rid.insert(gt_rid.value(i), gt_mt.value(i));
2448 }
2449
2450 let mut master_row_by_mid: HashMap<&str, usize> = HashMap::new();
2457 for i in 0..ds.num_rows() {
2458 master_row_by_mid.entry(ds_mid.value(i)).or_insert(i);
2459 }
2460
2461 let compare_cols: Vec<usize> = (0..ds.num_columns())
2462 .filter(|&i| {
2463 let field = ds.schema().field(i).clone();
2464 matches!(field.data_type(), DataType::Utf8)
2465 && !matches!(
2466 field.name().as_str(),
2467 "record_id" | "domain" | "entity_type" | "master_id"
2468 )
2469 && !field.name().ends_with("_id")
2470 })
2471 .collect();
2472
2473 let mut max_changed_cols = 0usize;
2474 for i in 0..ds.num_rows() {
2475 if match_type_by_rid.get(ds_rid.value(i)) != Some(&"fuzzy_dup") {
2476 continue;
2477 }
2478 let Some(&master_i) = master_row_by_mid.get(ds_mid.value(i)) else {
2479 continue;
2480 };
2481 let mut changed = 0usize;
2482 for &col_i in &compare_cols {
2483 let col = ds.column(col_i).as_string::<i32>();
2484 let a = (!col.is_null(i)).then(|| col.value(i));
2485 let b = (!col.is_null(master_i)).then(|| col.value(master_i));
2486 if a != b {
2487 changed += 1;
2488 }
2489 }
2490 max_changed_cols = max_changed_cols.max(changed);
2491 }
2492
2493 assert!(
2494 max_changed_cols >= 2,
2495 "expected at least one fuzzy_dup row to differ from its master on \
2496 2+ columns (evidence of multiple noise passes), got max {max_changed_cols}"
2497 );
2498
2499 let _ = std::fs::remove_dir_all(&dir);
2500 }
2501
2502 #[test]
2503 fn test_record_id_and_pad_string_roundtrip() {
2504 let rid = record_id_string(7);
2505 assert_eq!(rid, "R-0000000000007");
2506 let pad = pad_string(7);
2507 assert_eq!(pad, "0000000000007");
2508 let strs = record_id_strs(0..3);
2509 assert_eq!(strs.len(), 3);
2510 assert_eq!(strs[0], record_id_string(0));
2511 assert_eq!(strs[2], record_id_string(2));
2512 }
2513
2514 #[test]
2515 fn test_parse_record_idx_roundtrip() {
2516 for i in [0usize, 7, 1_234_567, 300_000_000] {
2517 assert_eq!(parse_record_idx(&record_id_string(i)), Some(i as u64));
2518 }
2519 assert_eq!(parse_record_idx("garbage"), None);
2520 assert_eq!(parse_record_idx("R-1"), None);
2521 assert_eq!(parse_record_idx("X-0000000000007"), None);
2522 }
2523
2524 #[test]
2525 fn test_pack_master_key_roundtrip_and_ordering() {
2526 let mid = |e: usize, i: usize| format!("{}-{}", entity_prefix(e), pad_string(i));
2527 assert!(pack_master_key(&mid(0, 0)).is_some());
2528 assert!(pack_master_key(&mid(0, 1)).unwrap() > pack_master_key(&mid(0, 0)).unwrap());
2530 assert!(
2533 pack_master_key(&mid(1, 0)).unwrap()
2534 > pack_master_key(&mid(0, 9_999_999_999_999)).unwrap()
2535 );
2536 assert_eq!(pack_master_key("garbage"), None);
2537 assert_eq!(pack_master_key("HN-0000000000007"), None);
2538 assert_eq!(pack_master_key("CANARY-000-0-abc"), None);
2539 }
2540
2541 #[test]
2551 #[ignore]
2552 fn hunt_h1_zip_cost_vs_column_count() {
2553 use arrow::array::StringArray;
2554 use arrow::datatypes::{DataType, Field, Schema};
2555
2556 fn make_batch(n_rows: usize, n_cols: usize) -> RecordBatch {
2557 let fields: Vec<Field> = (0..n_cols)
2558 .map(|i| Field::new(format!("col{i}"), DataType::Utf8, false))
2559 .collect();
2560 let cols: Vec<ArrayRef> = (0..n_cols)
2561 .map(|_| {
2562 Arc::new(StringArray::from(
2563 (0..n_rows)
2564 .map(|_| "abcdefghijklmnop".to_string())
2565 .collect::<Vec<_>>(),
2566 )) as ArrayRef
2567 })
2568 .collect();
2569 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
2570 }
2571
2572 const N_ROWS: usize = 100_000;
2573 const N_REPS: usize = 5;
2574 const EXTRA_PASSES: usize = 2;
2578
2579 eprintln!("\n[hunt_h1] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
2580 eprintln!("[hunt_h1] C=total columns, T=1 noised column (fixed) -- H1 predicts cost ∝ C");
2581
2582 for &n_cols in &[8usize, 15, 23] {
2583 let plan_cols = vec!["col0".to_string()];
2584 let all_types: Vec<&str> = vec!["typo"];
2585 let mut total = std::time::Duration::ZERO;
2586
2587 for rep in 0..N_REPS {
2588 let batch = make_batch(N_ROWS, n_cols);
2589 let mut rng = Rng::new(42 + rep as u64);
2590 let t0 = std::time::Instant::now();
2591
2592 let (noisy, mut unchanged, _) =
2593 apply_noise_with_retry(&batch, "typo", &plan_cols, &mut rng, &[]).unwrap();
2594 let mut current = noisy;
2595 for _ in 0..EXTRA_PASSES {
2596 current = apply_extra_pass_per_row(
2597 ¤t,
2598 &mut unchanged,
2599 &all_types,
2600 &mut rng,
2601 &[],
2602 )
2603 .unwrap();
2604 }
2605 std::hint::black_box(¤t);
2606 total += t0.elapsed();
2607 }
2608
2609 let avg = total / N_REPS as u32;
2610 eprintln!(
2611 "[hunt_h1] C={n_cols:2} avg={:>8.3}ms per_column={:>7.4}ms per_row_per_column={:>10.6}us",
2612 avg.as_secs_f64() * 1000.0,
2613 avg.as_secs_f64() * 1000.0 / n_cols as f64,
2614 avg.as_secs_f64() * 1_000_000.0 / (N_ROWS * n_cols) as f64,
2615 );
2616 }
2617 }
2618
2619 #[test]
2636 #[ignore]
2637 fn hunt_h1_residual_kyc_vs_nonprofit() {
2638 use arrow::array::StringArray;
2639 use arrow::datatypes::{DataType, Field, Schema};
2640
2641 fn make_batch(n_rows: usize, col_names: &[&str]) -> RecordBatch {
2642 let fields: Vec<Field> = col_names
2643 .iter()
2644 .map(|n| Field::new(*n, DataType::Utf8, true))
2645 .collect();
2646 let cols: Vec<ArrayRef> = col_names
2647 .iter()
2648 .map(|_| {
2649 Arc::new(StringArray::from(
2650 (0..n_rows)
2651 .map(|_| "abcdefghijklmnop testword".to_string())
2652 .collect::<Vec<_>>(),
2653 )) as ArrayRef
2654 })
2655 .collect();
2656 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
2657 }
2658
2659 const N_ROWS: usize = 100_000;
2660 const N_REPS: usize = 5;
2661 const HELL_TYPES: &[&str] = &[
2662 "typo_aggressive",
2663 "visual",
2664 "unicode_pollution",
2665 "names",
2666 "dates",
2667 "identifiers",
2668 "addresses",
2669 "companies",
2670 "extra",
2671 ];
2672 const HELL_PASSES: usize = 8; let kyc_cols: &[&str] = &[
2675 "party_id",
2676 "given_name",
2677 "family_name",
2678 "middle_name",
2679 "birth_date",
2680 "birth_place",
2681 "nationality",
2682 "email_address",
2683 "mobile_phone",
2684 "personal_administrative_number",
2685 "tax_id",
2686 "residential_address_line1",
2687 "residential_city",
2688 "residential_postal_code",
2689 "residential_state",
2690 "residential_country",
2691 "occupation",
2692 "risk_score",
2693 "document_type",
2694 "document_number",
2695 "source_system",
2696 ];
2697 let nonprofit_cols: &[&str] = &[
2698 "donation_id",
2699 "donor_id",
2700 "campaign_id",
2701 "donation_date",
2702 "payment_method",
2703 "source_system",
2704 ];
2705
2706 eprintln!("\n[hunt_h1_residual] N_ROWS={N_ROWS} N_REPS={N_REPS} HELL_PASSES={HELL_PASSES}");
2707
2708 {
2713 let batch = make_batch(N_ROWS, kyc_cols);
2714 eprintln!("[hunt_h1_residual] per-category breakdown on kyc.natural_person:");
2715 for &ttype in HELL_TYPES {
2716 let mut rng = Rng::new(42);
2717 let t0 = std::time::Instant::now();
2718 let (_, _, tcols) =
2719 apply_noise_with_retry(&batch, ttype, &[], &mut rng, &[]).unwrap();
2720 eprintln!(
2721 "[hunt_h1_residual] {ttype:20} {:>8.3}ms target_cols={:?}",
2722 t0.elapsed().as_secs_f64() * 1000.0,
2723 tcols,
2724 );
2725 }
2726 }
2727
2728 for (label, col_names) in [
2729 ("kyc.natural_person", kyc_cols),
2730 ("nonprofit.donation", nonprofit_cols),
2731 ] {
2732 let mut total = std::time::Duration::ZERO;
2733 let mut total_target_hits = 0usize;
2734
2735 for rep in 0..N_REPS {
2736 let batch = make_batch(N_ROWS, col_names);
2737 let mut rng = Rng::new(42 + rep as u64);
2738 let t0 = std::time::Instant::now();
2739
2740 let (noisy, mut unchanged, t0cols) =
2743 apply_noise_with_retry(&batch, "identifiers", &[], &mut rng, &[]).unwrap();
2744 total_target_hits += t0cols.len();
2745 let mut current = noisy;
2746 for _ in 1..HELL_PASSES {
2747 current = apply_extra_pass_per_row(
2748 ¤t,
2749 &mut unchanged,
2750 HELL_TYPES,
2751 &mut rng,
2752 &[],
2753 )
2754 .unwrap();
2755 }
2756 std::hint::black_box(¤t);
2757 total += t0.elapsed();
2758 }
2759
2760 let avg = total / N_REPS as u32;
2761 eprintln!(
2762 "[hunt_h1_residual] {label:20} C={:2} avg={:>8.3}ms avg_first_pass_target_cols={:.1}",
2763 col_names.len(),
2764 avg.as_secs_f64() * 1000.0,
2765 total_target_hits as f64 / N_REPS as f64,
2766 );
2767 }
2768 }
2769}