1use std::collections::HashMap;
7use std::sync::Arc;
8
9use arrow::array::{
10 Array, ArrayBuilder, ArrayRef, AsArray, StringArray, StringBuilder, UInt64Array, UInt64Builder,
11};
12use arrow::datatypes::{DataType, Field, Schema};
13use arrow::record_batch::RecordBatch;
14use rayon::prelude::*;
15use rustc_hash::FxHashMap;
16use serde::Deserialize;
17
18use crate::context::Context;
19use crate::rng::Rng;
20
21#[derive(Debug, Deserialize)]
24pub struct PipelineConfig {
25 pub domain: String,
26 pub size: usize,
27 pub seed: u64,
28 pub difficulty: String,
29 pub output_format: String,
30 pub run_id: String,
31
32 pub entity_plans: Vec<EntityPlan>,
33 pub hard_neg_types: Vec<HnTypeConfig>,
34
35 pub hard_neg_ratio: f64,
36
37 pub noise_passes: usize,
42
43 pub graph_enabled: bool,
44 pub graph_format: String,
45}
46
47#[derive(Debug, Deserialize)]
48pub struct EntityPlan {
49 pub name: String,
50 pub n_base: usize,
51 pub identifier_col: Option<String>,
52 pub columns_json: String,
53 pub noise_types: Vec<NoisePlanEntry>,
54 pub fk_remaps: Vec<FkRemapEntry>,
55 #[serde(default)]
62 pub pool_only: bool,
63}
64
65pub(crate) const FK_POOL_CAP: usize = 200_000;
72
73#[derive(Debug, Deserialize)]
74pub struct NoisePlanEntry {
75 pub noise_type: String,
76 pub columns: Vec<String>,
77 pub count: usize,
78}
79
80#[derive(Debug, Deserialize)]
81pub struct FkRemapEntry {
82 pub source_col: String,
83 pub target_entity: String,
84 #[serde(default = "default_fk_weight")]
88 pub fk_weight: f64,
89}
90
91fn default_fk_weight() -> f64 {
92 1.0
93}
94
95#[derive(Debug, Deserialize)]
96pub struct HnTypeConfig {
97 pub entity_type: String,
98 pub config_json: String,
99 pub count: usize,
100}
101
102#[derive(Debug)]
103pub struct PipelineOutput {
104 pub output_files: Vec<String>,
105 pub gt_file: String,
106 pub stats: PipelineStats,
107 pub nodes: Option<String>,
108 pub edges: Option<String>,
109}
110
111#[derive(Debug)]
112pub struct PipelineStats {
113 pub total_records: usize,
114 pub exact_dups: usize,
117 pub fuzzy_dups: usize,
120 pub hard_negs: usize,
121 pub uniques: usize,
122 pub masters: usize,
123}
124
125#[derive(Debug, Clone, Default)]
132pub struct ChunkOffsets {
133 pub rid_offset: usize,
135 pub master_local: Vec<usize>,
141 pub hn_master_counter: u64,
143}
144
145impl ChunkOffsets {
146 fn master_base(&self, entity_idx: usize) -> usize {
147 self.master_local.get(entity_idx).copied().unwrap_or(0)
148 }
149}
150
151pub(crate) const RID_LEN: usize = 15; pub(crate) const PAD_LEN: usize = 13; #[cfg(test)]
173#[inline]
174pub(crate) fn record_id_string(i: usize) -> String {
175 let mut buf = [0u8; RID_LEN];
176 buf[0] = b'R';
177 buf[1] = b'-';
178 let mut n = i;
179 for j in (2..RID_LEN).rev() {
180 buf[j] = b'0' + (n % 10) as u8;
181 n /= 10;
182 }
183 String::from_utf8(buf.to_vec()).expect("record_id_string: buffer is 'R', '-' and ASCII digits")
184}
185#[cfg(test)]
186#[inline]
187pub(crate) fn pad_string(i: usize) -> String {
188 let mut buf = [0u8; PAD_LEN];
189 let mut n = i;
190 for j in (0..PAD_LEN).rev() {
191 buf[j] = b'0' + (n % 10) as u8;
192 n /= 10;
193 }
194 String::from_utf8(buf.to_vec()).expect("pad_string: buffer is always ASCII digits")
195}
196
197#[cfg(test)]
198pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
199 range.map(record_id_string).collect()
200}
201
202#[inline]
208pub(crate) fn write_digits(buf: &mut Vec<u8>, val: u64, width: usize) {
209 let start = buf.len();
210 buf.resize(start + width, b'0');
211 let mut n = val;
212 for j in (0..width).rev() {
213 buf[start + j] = b'0' + (n % 10) as u8;
214 n /= 10;
215 }
216}
217
218#[inline]
224pub(crate) fn append_record_id(buf: &mut Vec<u8>, idx: u64) {
225 buf.push(b'R');
226 buf.push(b'-');
227 write_digits(buf, idx, PAD_LEN);
228}
229
230pub(crate) fn record_id_array(range: std::ops::Range<usize>) -> ArrayRef {
235 let n = range.len();
236 let mut i = range.start as u64;
237 crate::buf_gen::build_string_array(n, RID_LEN, |buf| {
238 buf.push(b'R');
239 buf.push(b'-');
240 write_digits(buf, i, PAD_LEN);
241 i += 1;
242 })
243}
244
245fn master_id_array_range(prefix: &str, start: usize, n: usize) -> ArrayRef {
249 let width = prefix.len() + 1 + PAD_LEN;
250 let mut i = start as u64;
251 crate::buf_gen::build_string_array(n, width, |buf| {
252 buf.extend_from_slice(prefix.as_bytes());
253 buf.push(b'-');
254 write_digits(buf, i, PAD_LEN);
255 i += 1;
256 })
257}
258
259fn master_id_array_from_indices(
264 prefix: &str,
265 mut indices: impl Iterator<Item = usize>,
266 n: usize,
267) -> ArrayRef {
268 let width = prefix.len() + 1 + PAD_LEN;
269 crate::buf_gen::build_string_array(n, width, |buf| {
270 let idx = indices
271 .next()
272 .expect("master_id_array_from_indices: fewer than n indices");
273 buf.extend_from_slice(prefix.as_bytes());
274 buf.push(b'-');
275 write_digits(buf, idx as u64, PAD_LEN);
276 })
277}
278
279fn hn_master_id_array(start: u64, n: usize) -> ArrayRef {
283 const HN_WIDTH: usize = 12; let mut id = start;
285 crate::buf_gen::build_string_array(n, HN_WIDTH, |buf| {
286 buf.extend_from_slice(b"HN-");
287 write_digits(buf, id, 9);
288 id += 1;
289 })
290}
291
292pub(crate) fn parse_record_idx(rid: &str) -> Option<u64> {
298 if rid.len() != RID_LEN || !rid.starts_with("R-") {
299 return None;
300 }
301 rid[2..].parse().ok()
302}
303
304pub(crate) fn pack_master_key(mid: &str) -> Option<u64> {
318 let bytes = mid.as_bytes();
319 if bytes.len() != 1 + 5 + 1 + PAD_LEN || bytes[0] != b'E' || bytes[6] != b'-' {
320 return None;
321 }
322 let entity_idx: u64 = mid[1..6].parse().ok()?;
323 let local_idx: u64 = mid[7..7 + PAD_LEN].parse().ok()?;
324 Some((entity_idx << 44) | local_idx)
325}
326
327fn distribute_by_weight(total: usize, weights: &[usize]) -> Vec<usize> {
335 let weight_sum: usize = weights.iter().sum();
336 if weight_sum == 0 || total == 0 {
337 return vec![0; weights.len()];
338 }
339 let raw: Vec<f64> = weights
340 .iter()
341 .map(|&w| total as f64 * w as f64 / weight_sum as f64)
342 .collect();
343 let mut counts: Vec<usize> = raw.iter().map(|&r| r as usize).collect();
344 let assigned: usize = counts.iter().sum();
345 let mut remainder = total - assigned;
346 let mut order: Vec<usize> = (0..counts.len()).collect();
347 order.sort_by(|&a, &b| {
348 let fa = raw[a] - counts[a] as f64;
349 let fb = raw[b] - counts[b] as f64;
350 fb.partial_cmp(&fa).unwrap_or(std::cmp::Ordering::Equal)
351 });
352 for &i in &order {
353 if remainder == 0 {
354 break;
355 }
356 counts[i] += 1;
357 remainder -= 1;
358 }
359 counts
360}
361
362pub(crate) fn entity_prefix(entity_idx: usize) -> String {
371 format!("E{:05}", entity_idx)
372}
373
374type FkPoolMap = HashMap<String, RecordBatch>;
377
378struct HnPool {
379 batch: RecordBatch,
380 total_count: usize,
381 record_ids: ArrayRef,
384}
385
386const PERSON_NAME_WORDS: &[&str] = &[
393 "name",
394 "first_name",
402 "last_name",
403 "given",
404 "family",
405 "middle",
406 "operator",
407 "technician",
408 "inspector",
409 "claimant",
410 "consignee",
411 "investigator",
412 "collected_by",
413 "performed_by",
414 "assigned_to",
415];
416
417const COMPANY_NAME_WORDS: &[&str] = &[
423 "company",
424 "legal",
425 "trading",
426 "name",
427 "supplier",
428 "manufacturer",
429 "sponsor",
430 "institution",
431 "journal",
432 "funder",
433 "employer",
434 "law_firm",
435 "affiliation",
436 "reported_by",
437 "buyer",
438 "shipper",
439 "airline",
440];
441
442fn is_specifically_person_name(lower: &str) -> bool {
456 PERSON_NAME_WORDS
457 .iter()
458 .any(|&w| w != "name" && lower.contains(w))
459}
460
461fn is_specifically_company_name(lower: &str) -> bool {
464 COMPANY_NAME_WORDS
465 .iter()
466 .any(|&w| w != "name" && lower.contains(w))
467}
468
469pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
475 let lower = col_name.to_lowercase();
476 match noise_type {
477 "typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
478 | "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
479 | "blocking_fail" => {
480 if lower.contains("email") || lower.contains("ip_address") {
485 return false;
486 }
487 contains_any(&lower, &["address", "street", "city"])
492 || contains_any(&lower, PERSON_NAME_WORDS)
493 || contains_any(&lower, COMPANY_NAME_WORDS)
494 }
495 "names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
496 contains_any(&lower, PERSON_NAME_WORDS) && !is_specifically_company_name(&lower)
502 }
503 "dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
504 contains_any(
520 &lower,
521 &[
522 "date",
523 "birth",
524 "incorporation",
525 "founding",
526 "time",
527 "since",
528 "period",
529 ],
530 )
531 }
532 "missing" | "missing_pattern" => {
533 contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
534 }
535 "identifiers"
536 | "corrupt_email"
537 | "corrupt_phone"
538 | "corrupt_national_id"
539 | "corrupt_siren"
540 | "national_id_corrupt"
541 | "phone_corrupt"
542 | "email_corrupt"
543 | "siren_corrupt" => contains_any(
544 &lower,
545 &[
546 "email",
547 "phone",
548 "ssn",
549 "pan",
550 "passport",
551 "account",
552 "medicare",
553 "license",
554 "siren",
555 "national_id",
556 "registration",
563 "sku",
573 "barcode",
574 "imei",
575 "imsi",
576 "iccid",
577 ],
578 ),
579 "extra"
580 | "name_null"
581 | "dob_null"
582 | "blocking_fail_initial"
583 | "blocking_fail_partial"
584 | "fuzzy_match"
585 | "phonetic" => {
586 if lower.contains("email") || lower.contains("ip_address") {
592 return false;
593 }
594 contains_any(
599 &lower,
600 &["address", "street", "city", "note", "comment", "bio"],
601 ) || contains_any(&lower, PERSON_NAME_WORDS)
602 || contains_any(&lower, COMPANY_NAME_WORDS)
603 }
604 "companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
605 contains_any(&lower, COMPANY_NAME_WORDS) && !is_specifically_person_name(&lower)
611 }
612 "addresses" | "address_scramble" | "postal_corrupt" => {
613 !lower.contains("ip_address")
636 && !lower.contains("email")
637 && contains_any(
638 &lower,
639 &[
640 "address", "street", "postal", "city", "state", "country", "location",
641 ],
642 )
643 }
644 "exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
645 | "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
646 | "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
647 | "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
648 | "combo_ultimate" | "french_address" => {
649 false
651 }
652 _ => true,
654 }
655}
656
657fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
658 let mut matched: Vec<String> = Vec::new();
659 for field in schema.fields() {
660 if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
661 continue;
662 }
663 if exclude_cols.contains(field.name()) {
664 continue;
665 }
666 if noise_type_targets_column(noise_type, field.name()) {
667 matched.push(field.name().clone());
668 }
669 }
670 if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
672 matched.truncate(2);
673 }
674 matched.sort();
675 matched.dedup();
676 matched
677}
678
679fn contains_any(s: &str, patterns: &[&str]) -> bool {
680 patterns.iter().any(|p| s.contains(p))
681}
682
683fn apply_noise_to_batch(
690 rb: &RecordBatch,
691 noise_type: &str,
692 plan_cols: &[String],
693 rng: &mut Rng,
694 exclude_cols: &[String],
695) -> Result<(RecordBatch, Vec<String>), String> {
696 let schema = rb.schema();
697 let n_cols = rb.num_columns();
698
699 let target_cols: Vec<String> = if plan_cols.is_empty() {
701 match_noise_columns(&schema, noise_type, exclude_cols)
702 } else {
703 plan_cols.to_vec()
704 };
705
706 let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();
708
709 let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
710 new_columns[idx]
711 .clone()
712 .unwrap_or_else(|| rb.column(idx).clone())
713 };
714
715 match noise_type {
716 "blocking_fail" => {
717 for col_name in &target_cols {
718 let lower = col_name.to_lowercase();
719 if let Ok(idx) = schema.index_of(col_name) {
720 let col = get_col(&new_columns, idx, rb);
721 if lower.contains("first") || lower.contains("given") {
722 new_columns[idx] =
723 Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
724 } else {
725 new_columns[idx] =
726 Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
727 }
728 }
729 }
730 }
731 "swap" | "full_swap" => {
732 if target_cols.len() >= 2 {
733 let name_a = &target_cols[0];
734 let name_b = &target_cols[1];
735 if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
736 let a = get_col(&new_columns, idx_a, rb);
737 let b = get_col(&new_columns, idx_b, rb);
738 new_columns[idx_a] = Some(b);
739 new_columns[idx_b] = Some(a);
740 }
741 }
742 }
743 "missing_pattern" => {
744 let mut phone_idx: Option<usize> = None;
745 let mut email_idx: Option<usize> = None;
746 for col_name in &target_cols {
747 let lower = col_name.to_lowercase();
748 if let Ok(idx) = schema.index_of(col_name) {
749 if lower.contains("phone") || lower.contains("mobile") {
750 phone_idx = Some(idx);
751 } else if lower.contains("email") {
752 email_idx = Some(idx);
753 }
754 }
755 }
756 if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
757 let phone_arr = get_col(&new_columns, pi, rb);
758 let email_arr = get_col(&new_columns, ei, rb);
759 let phone_s = phone_arr.as_string::<i32>();
760 let email_s = email_arr.as_string::<i32>();
761 let n = phone_s.len();
762 let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
763 for i in 0..n {
764 let pv = if phone_s.is_null(i) {
765 ""
766 } else {
767 phone_s.value(i)
768 };
769 let ev = if email_s.is_null(i) {
770 ""
771 } else {
772 email_s.value(i)
773 };
774 if !ev.is_empty() && rng.next_usize(2) == 0 {
775 builder.append_value("");
776 } else {
777 builder.append_value(pv);
778 }
779 }
780 new_columns[pi] = Some(Arc::new(builder.finish()));
781 }
782 }
783 _ => {
784 for col_name in &target_cols {
785 if let Ok(col_idx) = schema.index_of(col_name) {
786 let col = get_col(&new_columns, col_idx, rb);
787 let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
788 .map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
789 new_columns[col_idx] = Some(result);
790 }
791 }
792 }
793 }
794
795 let final_columns: Vec<ArrayRef> = (0..n_cols)
797 .map(|i| {
798 new_columns[i]
799 .clone()
800 .unwrap_or_else(|| rb.column(i).clone())
801 })
802 .collect();
803
804 let batch =
805 RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))?;
806 Ok((batch, target_cols))
807}
808
809fn unchanged_row_mask(
820 orig: &RecordBatch,
821 noised: &RecordBatch,
822 target_cols: &[String],
823) -> Result<Vec<bool>, String> {
824 let n = orig.num_rows();
825 unchanged_row_mask_within(orig, noised, target_cols, &vec![true; n])
826}
827
828fn unchanged_row_mask_within(
839 orig: &RecordBatch,
840 noised: &RecordBatch,
841 target_cols: &[String],
842 only: &[bool],
843) -> Result<Vec<bool>, String> {
844 let n = orig.num_rows();
845 let mut changed: Vec<bool> = only.iter().map(|&keep| !keep).collect();
848 let schema = orig.schema();
849 for col_name in target_cols {
850 let Ok(idx) = schema.index_of(col_name) else {
851 continue;
852 };
853 let a = orig.column(idx).as_string::<i32>();
854 let b = noised.column(idx).as_string::<i32>();
855 #[allow(clippy::needless_range_loop)]
858 for i in 0..n {
859 if changed[i] {
860 continue;
861 }
862 let av = (!a.is_null(i)).then(|| a.value(i));
863 let bv = (!b.is_null(i)).then(|| b.value(i));
864 if av != bv {
865 changed[i] = true;
866 }
867 }
868 }
869 for c in changed.iter_mut() {
870 *c = !*c;
871 }
872 Ok(changed)
873}
874
875fn apply_noise_with_retry(
898 orig: &RecordBatch,
899 noise_type: &str,
900 plan_cols: &[String],
901 rng: &mut Rng,
902 exclude_cols: &[String],
903) -> Result<(RecordBatch, Vec<bool>, Vec<String>), String> {
904 const MAX_NOISE_ATTEMPTS: usize = 4;
905 let (mut noisy, target_cols) =
906 apply_noise_to_batch(orig, noise_type, plan_cols, rng, exclude_cols)?;
907 if target_cols.is_empty() {
908 return Ok((noisy, vec![true; orig.num_rows()], target_cols));
909 }
910 let schema = noisy.schema();
911 let target_idxs: Vec<usize> = target_cols
912 .iter()
913 .filter_map(|c| schema.index_of(c).ok())
914 .collect();
915 let mut unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
916 let mut attempt = 1;
917 while unchanged.iter().any(|&u| u) && attempt < MAX_NOISE_ATTEMPTS {
918 let mut retry_rng = Rng::new(rng.next_u64());
919 let (retried, _) =
920 apply_noise_to_batch(orig, noise_type, plan_cols, &mut retry_rng, exclude_cols)?;
921 let mask = arrow::array::BooleanArray::from_iter(unchanged.iter().map(|&u| Some(u)));
922 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
923 for &i in &target_idxs {
924 merged_cols[i] =
925 arrow::compute::kernels::zip::zip(&mask, retried.column(i), noisy.column(i))
926 .map_err(|e| format!("zip noise retry col {i}: {e}"))?;
927 }
928 noisy = RecordBatch::try_new(noisy.schema(), merged_cols)
929 .map_err(|e| format!("rebuild retried dup batch: {e}"))?;
930 unchanged = unchanged_row_mask_within(orig, &noisy, &target_cols, &unchanged)?;
935 attempt += 1;
936 }
937 Ok((noisy, unchanged, target_cols))
938}
939
940fn apply_extra_pass_per_row(
967 noisy: &RecordBatch,
968 unchanged: &mut [bool],
969 all_types: &[&str],
970 rng: &mut Rng,
971 exclude_cols: &[String],
972) -> Result<RecordBatch, String> {
973 let n = noisy.num_rows();
974 if all_types.is_empty() || n == 0 {
975 return Ok(noisy.clone());
976 }
977 let per_row_type: Vec<usize> = (0..n).map(|_| rng.next_usize(all_types.len())).collect();
978
979 let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
998 let mut row_group: Vec<u32> = Vec::with_capacity(n);
999 let mut row_pos: Vec<u32> = Vec::with_capacity(n);
1000 for (i, &t) in per_row_type.iter().enumerate() {
1001 row_pos.push(groups[t].len() as u32);
1002 row_group.push(t as u32);
1003 groups[t].push(i as u32);
1004 }
1005
1006 let schema = noisy.schema();
1007 let mut pass_unchanged = vec![true; n];
1008 let mut noised_by_group: Vec<Option<RecordBatch>> = vec![None; all_types.len()];
1009 let mut group_target_idxs: Vec<Vec<usize>> = vec![Vec::new(); all_types.len()];
1015 let mut touched_cols: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
1016
1017 for (type_idx, idxs) in groups.into_iter().enumerate() {
1018 if idxs.is_empty() {
1019 continue;
1020 }
1021 let ttype = all_types[type_idx];
1022 let idx_arr = UInt64Array::from_iter_values(idxs.iter().map(|&i| i as u64));
1023 let subset = pick_rows(noisy, &idx_arr)?;
1024 let (noised_subset, subset_unchanged, target_cols) =
1025 apply_noise_with_retry(&subset, ttype, &[], rng, exclude_cols)?;
1026 let target_idxs: Vec<usize> = target_cols
1027 .iter()
1028 .filter_map(|c| schema.index_of(c).ok())
1029 .collect();
1030 touched_cols.extend(target_idxs.iter().copied());
1031 group_target_idxs[type_idx] = target_idxs;
1032 for (local_pos, &orig_i) in idxs.iter().enumerate() {
1033 pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
1034 }
1035 noised_by_group[type_idx] = Some(noised_subset);
1036 }
1037
1038 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
1039 for col_idx in touched_cols {
1040 let orig_col = noisy.column(col_idx).as_string::<i32>();
1041 let cap_bytes = orig_col.value_data().len().max(n);
1047 let mut builder = StringBuilder::with_capacity(n, cap_bytes);
1048 let per_group: Vec<Option<&arrow::array::StringArray>> = (0..group_target_idxs.len())
1057 .map(|g| {
1058 if group_target_idxs[g].contains(&col_idx) {
1059 noised_by_group[g]
1060 .as_ref()
1061 .map(|b| b.column(col_idx).as_string::<i32>())
1062 } else {
1063 None
1064 }
1065 })
1066 .collect();
1067 for i in 0..n {
1068 let g = row_group[i] as usize;
1069 match per_group[g] {
1070 Some(noised_col) => {
1071 let pos = row_pos[i] as usize;
1072 if noised_col.is_null(pos) {
1073 builder.append_null();
1074 } else {
1075 builder.append_value(noised_col.value(pos));
1076 }
1077 }
1078 None => {
1079 if orig_col.is_null(i) {
1080 builder.append_null();
1081 } else {
1082 builder.append_value(orig_col.value(i));
1083 }
1084 }
1085 }
1086 }
1087 merged_cols[col_idx] = Arc::new(builder.finish());
1088 }
1089 let current = RecordBatch::try_new(noisy.schema(), merged_cols)
1090 .map_err(|e| format!("rebuild extra-pass batch: {e}"))?;
1091
1092 for (u, pu) in unchanged.iter_mut().zip(pass_unchanged.iter()) {
1093 *u = *u && *pu;
1094 }
1095 Ok(current)
1096}
1097
1098fn topological_sort(
1102 plans: &[EntityPlan],
1103 name_to_idx: &HashMap<&str, usize>,
1104) -> Result<Vec<usize>, String> {
1105 let n = plans.len();
1106 let mut in_degree = vec![0usize; n];
1107 let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
1108 for i in 0..n {
1109 for remap in &plans[i].fk_remaps {
1110 let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
1111 continue;
1112 };
1113 if target == i {
1114 continue;
1115 }
1116 adj[target].push(i); in_degree[i] += 1;
1118 }
1119 }
1120 let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
1121 let mut order = Vec::with_capacity(n);
1122 let mut idx = 0;
1123 while idx < queue.len() {
1124 let v = queue[idx];
1125 idx += 1;
1126 order.push(v);
1127 for &next in &adj[v] {
1128 in_degree[next] -= 1;
1129 if in_degree[next] == 0 {
1130 queue.push(next);
1131 }
1132 }
1133 }
1134 if order.len() != n {
1135 return Err("FK dependency cycle detected among entities".into());
1136 }
1137 Ok(order)
1138}
1139
1140pub(crate) enum DatasetWriter {
1144 Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
1145 Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
1146}
1147
1148impl DatasetWriter {
1149 fn new(
1150 output_format: &str,
1151 path: &str,
1152 schema: &Arc<Schema>,
1153 metadata: &HashMap<String, String>,
1154 ) -> Result<Self, String> {
1155 let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
1156 if output_format == "parquet" {
1157 use parquet::basic::{Compression, ZstdLevel};
1158 use parquet::file::properties::WriterProperties;
1159 let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
1160 let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
1161 .iter()
1162 .map(|(k, v)| parquet::file::metadata::KeyValue {
1163 key: k.clone(),
1164 value: Some(v.clone()),
1165 })
1166 .collect();
1167 let props = WriterProperties::builder()
1168 .set_compression(Compression::ZSTD(zstd))
1169 .set_max_row_group_row_count(Some(1_000_000))
1173 .set_key_value_metadata(Some(meta_kv))
1174 .build();
1175 let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
1176 .map_err(|e| format!("ArrowWriter error: {e}"))?;
1177 Ok(DatasetWriter::Parquet(Box::new(writer)))
1178 } else {
1179 let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
1180 .map_err(|e| format!("IPC FileWriter error: {e}"))?;
1181 Ok(DatasetWriter::Ipc(Box::new(writer)))
1182 }
1183 }
1184
1185 pub(crate) fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
1186 match self {
1187 DatasetWriter::Ipc(w) => w.write(batch).map_err(|e| format!("write ipc: {e}")),
1188 DatasetWriter::Parquet(w) => w.write(batch).map_err(|e| format!("write parquet: {e}")),
1189 }
1190 }
1191
1192 fn finish(self) -> Result<(), String> {
1193 match self {
1194 DatasetWriter::Ipc(mut w) => w.finish().map_err(|e| format!("finish ipc: {e}")),
1195 DatasetWriter::Parquet(w) => w
1196 .close()
1197 .map(|_| ())
1198 .map_err(|e| format!("close parquet: {e}")),
1199 }
1200 }
1201}
1202
1203fn log_rss(label: &str) {
1215 if !log::log_enabled!(log::Level::Debug) {
1216 return;
1217 }
1218 let pid = sysinfo::Pid::from_u32(std::process::id());
1219 let mut sys = sysinfo::System::new();
1220 sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
1221 if let Some(p) = sys.process(pid) {
1222 log::debug!("[hunt_rss] {label}: rss={:.1} Mo", p.memory() as f64 / 1e6);
1223 }
1224}
1225
1226pub fn run_pipeline(
1229 ctx: &Context,
1230 config: &PipelineConfig,
1231 output_dir: &str,
1232) -> Result<PipelineOutput, String> {
1233 run_pipeline_with_progress(ctx, config, output_dir, None)
1234}
1235
1236pub fn run_pipeline_with_progress(
1243 ctx: &Context,
1244 config: &PipelineConfig,
1245 output_dir: &str,
1246 progress: Option<&mut dyn FnMut(usize, usize)>,
1247) -> Result<PipelineOutput, String> {
1248 run_pipeline_chunked(ctx, config, output_dir, progress, None).map(|(out, _)| out)
1249}
1250
1251pub fn run_pipeline_chunked(
1258 ctx: &Context,
1259 config: &PipelineConfig,
1260 output_dir: &str,
1261 mut progress: Option<&mut dyn FnMut(usize, usize)>,
1262 offsets: Option<&ChunkOffsets>,
1263) -> Result<(PipelineOutput, ChunkOffsets), String> {
1264 let offsets = offsets.cloned().unwrap_or_default();
1265 std::fs::create_dir_all(output_dir)
1266 .map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
1267 let t_start = std::time::Instant::now();
1268
1269 let t_alloc = t_start.elapsed().as_secs_f64();
1270 log_rss("start");
1271
1272 let t1e = std::time::Instant::now();
1276
1277 let metadata = build_metadata_map(config);
1283 let full_arc = Arc::new(build_full_schema(config, &metadata));
1285 let domain_dict = DictValues::new([config.domain.clone()]);
1293 let entity_type_dict = DictValues::new(
1294 config
1295 .entity_plans
1296 .iter()
1297 .filter(|p| !p.pool_only)
1298 .map(|p| p.name.clone()),
1299 );
1300 let dataset_ext = if config.output_format == "parquet" {
1301 "parquet"
1302 } else {
1303 "ipc"
1304 };
1305 let dataset_path = format!("{}/{}.{}", output_dir, config.run_id, dataset_ext);
1306 let mut writer =
1307 DatasetWriter::new(&config.output_format, &dataset_path, &full_arc, &metadata)?;
1308
1309 let name_to_idx: HashMap<&str, usize> = config
1311 .entity_plans
1312 .iter()
1313 .enumerate()
1314 .map(|(i, p)| (p.name.as_str(), i))
1315 .collect();
1316 let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;
1317
1318 let mut fk_pools: FkPoolMap = HashMap::new();
1320 let mut hn_pools: HashMap<String, HnPool> = HashMap::new();
1321
1322 let gt_ext = if config.output_format == "parquet" {
1326 "parquet"
1327 } else {
1328 "ipc"
1329 };
1330 let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
1331 let gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
1332 let mut gt_acc = crate::gt::GtAccumulator::new(>_draft_path)?;
1333
1334 let mut global_rid_offset = offsets.rid_offset;
1335
1336 const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
1338 let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
1339 let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
1340 let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
1341 let mut graph_nodes_path: Option<String> = None;
1342 let mut graph_edges_path: Option<String> = None;
1343 if config.graph_enabled {
1344 let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
1345 let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
1346 node_writer = Some(
1347 crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
1348 .map_err(|e| format!("init node writer: {e}"))?,
1349 );
1350 edge_writer = Some(
1351 crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata, subtype_dict(config))
1352 .map_err(|e| format!("init edge writer: {e}"))?,
1353 );
1354 graph_nodes_path = Some(nodes_ipc);
1355 graph_edges_path = Some(edges_ipc);
1356 }
1357
1358 let mut _t_write = 0.0f64;
1359 let mut _t_meta = 0.0f64;
1360 let mut _t_dup = 0.0f64;
1361 let mut _write_calls = 0u64;
1362 let hn_pool_sizing: std::collections::HashMap<&str, usize> =
1364 config
1365 .hard_neg_types
1366 .iter()
1367 .fold(std::collections::HashMap::new(), |mut m, hn| {
1368 let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
1369 *entry = (*entry).max(hn.count);
1370 m
1371 });
1372 let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
1373 let pool_only_names: std::collections::HashSet<&str> = config
1389 .entity_plans
1390 .iter()
1391 .filter(|p| p.pool_only)
1392 .map(|p| p.name.as_str())
1393 .collect();
1394
1395 let fk_targets: std::collections::HashSet<&str> = config
1396 .entity_plans
1397 .iter()
1398 .flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
1399 .collect();
1400
1401 for &plan_idx in &entity_order {
1402 let plan = &config.entity_plans[plan_idx];
1403 if plan.n_base == 0 {
1404 continue;
1405 }
1406
1407 let prefix = entity_prefix(plan_idx);
1408 let master_base = offsets.master_base(plan_idx);
1409 let parsed_columns = crate::entity_gen::parse_columns(&plan.columns_json)?;
1412
1413 let fk_targeted = fk_targets.contains(plan.name.as_str());
1415 let mut fk_builder = arrow::array::StringBuilder::new();
1416 let mut fk_rid_builder = arrow::array::StringBuilder::new();
1417 let mut fk_rid_buf: Vec<u8> = Vec::with_capacity(RID_LEN);
1421 let mut fk_count: usize = 0;
1422 let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
1427 let mut hn_slices: Vec<RecordBatch> = Vec::new();
1428 let mut hn_slice_rids: Vec<ArrayRef> = Vec::new();
1429 let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
1430 let mut fk_rng = Rng::new(config.seed.wrapping_add(42));
1431
1432 let batch_bounds: Vec<(usize, usize)> = (0..plan.n_base)
1438 .step_by(crate::entity_gen::BATCH_SIZE)
1439 .map(|offset| {
1440 (
1441 offset,
1442 (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE),
1443 )
1444 })
1445 .collect();
1446 let batch_weights: Vec<usize> = batch_bounds.iter().map(|&(_, bn)| bn).collect();
1447 let per_batch_noise_counts: Vec<Vec<usize>> = plan
1448 .noise_types
1449 .iter()
1450 .map(|n| distribute_by_weight(n.count, &batch_weights))
1451 .collect();
1452 let fk_exclude_cols: Vec<String> = plan
1458 .fk_remaps
1459 .iter()
1460 .map(|r| r.source_col.clone())
1461 .collect();
1462 let all_types: Vec<&str> = plan
1463 .noise_types
1464 .iter()
1465 .map(|n| n.noise_type.as_str())
1466 .collect();
1467
1468 let mut col_lookup: Option<Vec<Option<usize>>> = None;
1470 let mut null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1472 let mut const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1473
1474 for (batch_idx, offset) in (0..plan.n_base)
1476 .step_by(crate::entity_gen::BATCH_SIZE)
1477 .enumerate()
1478 {
1479 let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
1480 let batch_seed = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);
1481
1482 let rb = crate::entity_gen::generate_entity_batch_parsed(
1484 ctx,
1485 &parsed_columns,
1486 batch_n,
1487 batch_seed,
1488 offset,
1489 )?;
1490
1491 let mut fk_edges_by_remap: Vec<(String, ArrayRef, f64)> = Vec::new();
1493 let rb = if !plan.fk_remaps.is_empty() {
1494 let mut r = rb;
1495 for remap in &plan.fk_remaps {
1496 if let Some(pool) = fk_pools.get(&remap.target_entity) {
1497 let want_edge = config.graph_enabled
1501 && !pool_only_names.contains(remap.target_entity.as_str());
1502 let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
1503 &r,
1504 pool,
1505 &remap.source_col,
1506 &mut fk_rng,
1507 want_edge,
1508 )?;
1509 r = remapped;
1510 if want_edge && let Some(rids) = target_rids {
1511 fk_edges_by_remap.push((
1512 remap.source_col.clone(),
1513 rids,
1514 remap.fk_weight,
1515 ));
1516 }
1517 }
1518 }
1519 r
1520 } else {
1521 rb
1522 };
1523
1524 if fk_targeted
1527 && fk_count < FK_POOL_CAP
1528 && let Some(ref id_col) = plan.identifier_col
1529 && let Some(col) = rb.column_by_name(id_col)
1530 {
1531 let s = col.as_string::<i32>();
1532 for i in 0..s.len() {
1533 if fk_count >= FK_POOL_CAP {
1534 break;
1535 }
1536 if !s.is_null(i) {
1537 fk_builder.append_value(s.value(i));
1538 if config.graph_enabled {
1539 fk_rid_buf.clear();
1540 append_record_id(&mut fk_rid_buf, (global_rid_offset + i) as u64);
1541 fk_rid_builder.append_value(
1542 std::str::from_utf8(&fk_rid_buf)
1543 .expect("append_record_id writes only ASCII digits"),
1544 );
1545 }
1546 fk_count += 1;
1547 }
1548 }
1549 }
1550
1551 if plan.pool_only {
1557 continue;
1558 }
1559
1560 let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1562 if hn_targeted && hn_accum < hn_max_pool {
1563 let take = (hn_max_pool - hn_accum).min(batch_n);
1564 hn_slices.push(rb.slice(0, take));
1565 if config.graph_enabled {
1569 hn_slice_rids
1570 .push(record_id_array(global_rid_offset..global_rid_offset + take));
1571 }
1572 }
1573
1574 if col_lookup.is_none() {
1576 let plan_schema = rb.schema();
1577 col_lookup = Some(
1578 full_arc
1579 .fields()
1580 .iter()
1581 .skip(4)
1582 .map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1583 .collect(),
1584 );
1585 }
1586
1587 let batch_mid_arr = master_id_array_range(&prefix, master_base + offset, batch_n);
1589
1590 let rid_arr = record_id_array(global_rid_offset..global_rid_offset + batch_n);
1592 let t_m0 = std::time::Instant::now();
1593 let base_rb = add_metadata_and_align(
1594 &rb,
1595 &config.domain,
1596 &plan.name,
1597 &domain_dict,
1598 &entity_type_dict,
1599 &rid_arr,
1600 &batch_mid_arr,
1601 &full_arc,
1602 col_lookup.as_ref().unwrap(),
1603 &mut null_cache,
1604 &mut const_arr_cache,
1605 )?;
1606 _t_meta += t_m0.elapsed().as_secs_f64();
1607
1608 if config.graph_enabled
1610 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1611 {
1612 nw.write_batch(&base_rb)
1613 .map_err(|e| format!("write node: {e}"))?;
1614 for (subtype, target_rids, weight) in &fk_edges_by_remap {
1621 ew.push_batch(&rid_arr, target_rids, "fk", subtype, *weight)
1622 .map_err(|e| format!("push fk edge batch: {e}"))?;
1623 }
1624 }
1625
1626 let t_w0 = std::time::Instant::now();
1627 writer
1628 .write(&base_rb)
1629 .map_err(|e| format!("write base: {e}"))?;
1630 _t_write += t_w0.elapsed().as_secs_f64();
1631 _write_calls += 1;
1632
1633 gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;
1635
1636 global_rid_offset += batch_n;
1637 if let Some(cb) = &mut progress {
1638 cb(global_rid_offset, config.size);
1639 }
1640
1641 let has_dups: bool = per_batch_noise_counts
1649 .iter()
1650 .any(|counts| counts[batch_idx] > 0);
1651 if has_dups {
1652 let t_d0 = std::time::Instant::now();
1653 let mut dup_batches: Vec<RecordBatch> = Vec::new();
1654 let mut dup_master_idx_buf: Vec<usize> = Vec::new();
1655
1656 let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
1658 .noise_types
1659 .iter()
1660 .enumerate()
1661 .filter_map(|(ni, n)| {
1662 let cnt = per_batch_noise_counts[ni][batch_idx];
1663 if cnt == 0 {
1664 return None;
1665 }
1666 let mut b = UInt64Builder::with_capacity(cnt);
1667 for _ in 0..cnt {
1668 b.append_value(batch_rng.next_usize(batch_n) as u64);
1669 }
1670 Some((
1671 b.finish(),
1672 batch_rng.next_u64(),
1673 n.noise_type.as_str(),
1674 n.columns.as_slice(),
1675 cnt,
1676 ))
1677 })
1678 .collect();
1679
1680 let all_types_ref = &all_types;
1685 let noise_passes = config.noise_passes.max(1);
1686
1687 type DupNoiseResult = Result<(RecordBatch, Vec<usize>, Vec<bool>), String>;
1688 let fk_exclude_ref = &fk_exclude_cols;
1689 let rb_ref = &rb;
1690 let results: Vec<DupNoiseResult> = ndata
1703 .par_iter()
1704 .map(|(indices, seed, ntype, cols, cnt)| {
1705 let dup = pick_rows(rb_ref, indices)?;
1706 let mut rng = Rng::new(*seed);
1707
1708 let (mut noisy, mut unchanged, _) =
1714 apply_noise_with_retry(&dup, ntype, cols, &mut rng, fk_exclude_ref)?;
1715
1716 for _ in 1..noise_passes {
1728 noisy = apply_extra_pass_per_row(
1729 &noisy,
1730 &mut unchanged,
1731 all_types_ref,
1732 &mut rng,
1733 fk_exclude_ref,
1734 )?;
1735 }
1736
1737 let mut mb = Vec::with_capacity(*cnt);
1738 for j in 0..*cnt {
1739 let global_idx = master_base + offset + indices.value(j) as usize;
1744 mb.push(global_idx);
1745 }
1746 Ok((noisy, mb, unchanged))
1747 })
1748 .collect();
1749 let mut dup_is_identical_buf: Vec<bool> = Vec::new();
1750 for res in results {
1751 let (rb, mb, unchanged) = res?;
1752 dup_batches.push(rb);
1753 dup_master_idx_buf.extend(mb);
1754 dup_is_identical_buf.extend(unchanged);
1755 }
1756 _t_dup += t_d0.elapsed().as_secs_f64();
1757
1758 if !dup_batches.is_empty() {
1760 let dup_total = dup_master_idx_buf.len();
1761 let dup_rid_arr =
1762 record_id_array(global_rid_offset..global_rid_offset + dup_total);
1763 let dup_mid_arr = master_id_array_from_indices(
1764 &prefix,
1765 dup_master_idx_buf.iter().copied(),
1766 dup_total,
1767 );
1768
1769 let concated = if dup_batches.len() == 1 {
1770 dup_batches.into_iter().next().unwrap()
1771 } else {
1772 let schema = dup_batches[0].schema();
1773 let n_fields = dup_batches[0].num_columns();
1774 let mut concat_arrays = Vec::with_capacity(n_fields);
1775 for i in 0..n_fields {
1776 let refs: Vec<&dyn Array> =
1777 dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
1778 concat_arrays.push(
1779 arrow::compute::concat(&refs)
1780 .map_err(|e| format!("concat dup col {i}: {e}"))?,
1781 );
1782 }
1783 RecordBatch::try_new(schema, concat_arrays)
1784 .map_err(|e| format!("concat dups: {e}"))?
1785 };
1786
1787 let dup_rb_full = add_metadata_and_align(
1788 &concated,
1789 &config.domain,
1790 &plan.name,
1791 &domain_dict,
1792 &entity_type_dict,
1793 &dup_rid_arr,
1794 &dup_mid_arr,
1795 &full_arc,
1796 col_lookup.as_ref().unwrap(),
1797 &mut null_cache,
1798 &mut const_arr_cache,
1799 )?;
1800
1801 let t_wd = std::time::Instant::now();
1802 writer
1803 .write(&dup_rb_full)
1804 .map_err(|e| format!("write dups: {e}"))?;
1805 _t_write += t_wd.elapsed().as_secs_f64();
1806 _write_calls += 1;
1807
1808 if config.graph_enabled
1810 && let Some(nw) = node_writer.as_mut()
1811 {
1812 nw.write_batch(&dup_rb_full)
1813 .map_err(|e| format!("write dup node: {e}"))?;
1814 }
1815
1816 let is_identical_arr: ArrayRef =
1817 Arc::new(arrow::array::BooleanArray::from(dup_is_identical_buf));
1818 gt_acc.push_dup_batch(
1819 dup_rb_full.column(0),
1820 dup_rb_full.column(2),
1821 dup_rb_full.column(3),
1822 &is_identical_arr,
1823 )?;
1824 global_rid_offset += dup_total;
1825 }
1826 }
1827 }
1828
1829 if let Some(id_col) = &plan.identifier_col
1831 && fk_builder.len() > 0
1832 {
1833 let arr = Arc::new(fk_builder.finish()) as ArrayRef;
1834 if config.graph_enabled {
1835 let rid_arr = Arc::new(fk_rid_builder.finish()) as ArrayRef;
1839 let schema = Arc::new(Schema::new(vec![
1840 Field::new(id_col, DataType::Utf8, true),
1841 Field::new("record_id", DataType::Utf8, true),
1842 ]));
1843 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
1844 fk_pools.insert(plan.name.clone(), pool_rb);
1845 }
1846 } else {
1847 let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
1848 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
1849 fk_pools.insert(plan.name.clone(), pool_rb);
1850 }
1851 }
1852 }
1853
1854 if !hn_slices.is_empty() {
1856 let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1857 let batch = if hn_slices.len() == 1 {
1858 let only = hn_slices.into_iter().next().unwrap();
1865 let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
1866 pick_rows(&only, &idx)?
1867 } else {
1868 let schema = hn_slices[0].schema();
1869 let n_fields = schema.fields().len();
1870 let mut concat_arrays = Vec::with_capacity(n_fields);
1871 for i in 0..n_fields {
1872 let refs: Vec<&dyn Array> =
1873 hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
1874 concat_arrays.push(
1875 arrow::compute::concat(&refs)
1876 .map_err(|e| format!("hn pool concat col {i}: {e}"))?,
1877 );
1878 }
1879 RecordBatch::try_new(schema, concat_arrays)
1880 .map_err(|e| format!("hn pool concat batch: {e}"))?
1881 };
1882 let record_ids: ArrayRef = if config.graph_enabled {
1884 let refs: Vec<&dyn Array> = hn_slice_rids.iter().map(|a| a.as_ref()).collect();
1885 arrow::compute::concat(&refs).map_err(|e| format!("hn rid concat: {e}"))?
1886 } else {
1887 Arc::new(StringArray::new_null(0))
1888 };
1889 hn_pools.insert(
1890 plan.name.clone(),
1891 HnPool {
1892 batch,
1893 total_count,
1894 record_ids,
1895 },
1896 );
1897 }
1898
1899 log::debug!(
1909 "[hunt_gt_sets] after entity '{}' ({} base rows): dup_masters={} masters_with_exact_copy={}",
1910 plan.name,
1911 plan.n_base,
1912 gt_acc.dup_masters_len(),
1913 gt_acc.masters_with_exact_copy_len(),
1914 );
1915 log_rss(&format!(
1916 "after entity '{}' ({} base rows)",
1917 plan.name, plan.n_base
1918 ));
1919 }
1920 let t1e_elapsed = t1e.elapsed().as_secs_f64();
1921 log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
1922 log_rss("after entity batches (fk_pools+hn_pools built)");
1923
1924 let t2 = std::time::Instant::now();
1926 let mut hn_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1928 let mut hn_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1929 let mut hn_master_id_counter: u64 = offsets.hn_master_counter;
1937 for hn_cfg in config.hard_neg_types.iter() {
1938 if hn_cfg.count == 0 {
1939 continue;
1940 }
1941 let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
1942 continue;
1943 };
1944 let hn_total = pool_data.total_count;
1945
1946 let pool_rb = &pool_data.batch;
1947 let n_hn_max = hn_cfg.count.min(hn_total / 2);
1948 if n_hn_max == 0 {
1949 continue;
1950 }
1951
1952 let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
1953 pool_rb,
1954 &hn_cfg.config_json,
1955 n_hn_max,
1956 config.seed.wrapping_add(200),
1957 config.graph_enabled,
1958 )?;
1959
1960 let n_hn = hn_rb.num_rows();
1961 if n_hn == 0 {
1962 continue;
1963 }
1964
1965 let hn_rid_arr = record_id_array(global_rid_offset..global_rid_offset + n_hn);
1966 let hn_mid_arr = hn_master_id_array(hn_master_id_counter, n_hn);
1967 hn_master_id_counter += n_hn as u64;
1968
1969 let hn_schema = hn_rb.schema();
1971 let hn_col_lookup: Vec<Option<usize>> = full_arc
1972 .fields()
1973 .iter()
1974 .skip(4)
1975 .map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1976 .collect();
1977
1978 let hn_rb_full = add_metadata_and_align(
1979 &hn_rb,
1980 &config.domain,
1981 &hn_cfg.entity_type,
1982 &domain_dict,
1983 &entity_type_dict,
1984 &hn_rid_arr,
1985 &hn_mid_arr,
1986 &full_arc,
1987 &hn_col_lookup,
1988 &mut hn_null_cache,
1989 &mut hn_const_arr_cache,
1990 )?;
1991
1992 let t_wh = std::time::Instant::now();
1993 writer
1994 .write(&hn_rb_full)
1995 .map_err(|e| format!("write hn: {e}"))?;
1996 _t_write += t_wh.elapsed().as_secs_f64();
1997 _write_calls += 1;
1998
1999 if config.graph_enabled
2002 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
2003 {
2004 nw.write_batch(&hn_rb_full)
2005 .map_err(|e| format!("write hn node: {e}"))?;
2006 if let Some((idx_a, pattern)) = hn_src {
2007 let take_n = n_hn.min(idx_a.len());
2013 let idx_arr = arrow::array::UInt64Array::from_iter_values(
2014 idx_a.iter().take(take_n).map(|&i| i as u64),
2015 );
2016 let tgt_arr = arrow::compute::take(&pool_data.record_ids, &idx_arr, None)
2017 .map_err(|e| format!("take hn target ids: {e}"))?;
2018 let src_arr = hn_rid_arr.slice(0, take_n);
2019 ew.push_batch(&src_arr, &tgt_arr, "hard_neg", &pattern, 1.0)
2020 .map_err(|e| format!("push hn edge batch: {e}"))?;
2021 }
2022 }
2023
2024 gt_acc.push_other_batch(
2026 hn_rb_full.column(0),
2027 hn_rb_full.column(2),
2028 hn_rb_full.column(3),
2029 )?;
2030
2031 global_rid_offset += n_hn;
2032 }
2033 let t2_elapsed = t2.elapsed().as_secs_f64();
2034 drop(hn_pools);
2041 log_rss("after hard negatives");
2042
2043 {
2045 let mut canary_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
2046 let mut canary_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
2047 crate::canary::generate_all(
2048 ctx,
2049 config,
2050 &full_arc,
2051 &domain_dict,
2052 &entity_type_dict,
2053 &mut canary_null_cache,
2054 &mut canary_const_arr_cache,
2055 &mut global_rid_offset,
2056 &fk_pools,
2057 &mut writer,
2058 &mut node_writer,
2059 &mut gt_acc,
2060 )?;
2061 }
2062 log_rss("after canary records");
2063
2064 let t3 = std::time::Instant::now();
2066 writer
2067 .finish()
2068 .map_err(|e| format!("finish dataset writer: {e}"))?;
2069 let t3a_elapsed = t3.elapsed().as_secs_f64();
2070 log_rss("after dataset writer finish");
2071
2072 let t3b = std::time::Instant::now();
2073 let t_gt0 = std::time::Instant::now();
2074 let crate::gt::GtResult {
2075 n_exact_dup,
2076 n_fuzzy_dup,
2077 n_hard_neg,
2078 n_unique,
2079 n_masters,
2080 cluster_map,
2081 } = gt_acc.finish(
2082 config.difficulty.as_str(),
2083 &config.output_format,
2084 >_path,
2085 &metadata,
2086 config.graph_enabled,
2087 )?;
2088 let _t_gt_compute = 0.0f64;
2089 let _t_gt_write = t_gt0.elapsed().as_secs_f64();
2090 let t3b_elapsed = t3b.elapsed().as_secs_f64();
2091 let _t3_elapsed = t3.elapsed().as_secs_f64();
2092 log_rss("after GT finish (cluster_map built)");
2093
2094 if config.graph_enabled
2096 && let Some(ew) = edge_writer.as_mut()
2097 {
2098 crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
2099 .map_err(|e| format!("push dup clusters: {e}"))?;
2100 }
2101
2102 let duration = t_start.elapsed().as_secs_f64();
2103
2104 log::debug!(
2105 "[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",
2106 t_alloc,
2107 t1e_elapsed,
2108 t2_elapsed,
2109 t3a_elapsed,
2110 _t_meta,
2111 _t_dup,
2112 _t_write,
2113 _t_gt_compute,
2114 _t_gt_write,
2115 t3b_elapsed,
2116 duration,
2117 );
2118
2119 let stats = PipelineStats {
2120 total_records: global_rid_offset,
2121 exact_dups: n_exact_dup,
2122 fuzzy_dups: n_fuzzy_dup,
2123 hard_negs: n_hard_neg,
2124 uniques: n_unique,
2125 masters: n_masters,
2126 };
2127
2128 let mut graph_nodes_final: Option<String> = None;
2130 let mut graph_edges_final: Option<String> = None;
2131 if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
2132 nw.finish()
2133 .map_err(|e| format!("finish node writer: {e}"))?;
2134 ew.finish()
2135 .map_err(|e| format!("finish edge writer: {e}"))?;
2136 let mut np = graph_nodes_path.clone().unwrap();
2137 let mut ep = graph_edges_path.clone().unwrap();
2138 if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
2139 let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
2140 let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
2141 convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
2142 convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
2143 std::fs::remove_file(&np).ok();
2144 std::fs::remove_file(&ep).ok();
2145 np = np_pq;
2146 ep = ep_pq;
2147 }
2148 graph_nodes_final = Some(np);
2149 graph_edges_final = Some(ep);
2150 }
2151 log_rss("end (after graph finalize)");
2152
2153 let mut master_local = offsets.master_local.clone();
2157 if master_local.len() < config.entity_plans.len() {
2158 master_local.resize(config.entity_plans.len(), 0);
2159 }
2160 for (idx, plan) in config.entity_plans.iter().enumerate() {
2161 if !plan.pool_only {
2162 master_local[idx] = offsets.master_base(idx) + plan.n_base;
2163 }
2164 }
2165 let next_offsets = ChunkOffsets {
2166 rid_offset: global_rid_offset,
2167 master_local,
2168 hn_master_counter: hn_master_id_counter,
2169 };
2170
2171 Ok((
2172 PipelineOutput {
2173 output_files: vec![dataset_path],
2174 gt_file: gt_path,
2175 stats,
2176 nodes: graph_nodes_final,
2177 edges: graph_edges_final,
2178 },
2179 next_offsets,
2180 ))
2181}
2182
2183fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
2187 use arrow::ipc::reader::FileReader;
2188
2189 let ipc_file =
2190 std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
2191 let ipc_reader =
2192 FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
2193 let schema = ipc_reader.schema();
2194
2195 let parquet_file =
2196 std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
2197 let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
2198 let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
2199 .metadata()
2200 .iter()
2201 .map(|(k, v)| parquet::file::metadata::KeyValue {
2202 key: k.clone(),
2203 value: Some(v.clone()),
2204 })
2205 .collect();
2206 let props = parquet::file::properties::WriterProperties::builder()
2207 .set_compression(parquet::basic::Compression::ZSTD(zstd_level))
2208 .set_max_row_group_row_count(Some(1_000_000))
2209 .set_key_value_metadata(Some(meta_kv))
2210 .build();
2211 let mut parquet_writer =
2212 parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
2213 .map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;
2214
2215 for batch_result in ipc_reader {
2216 let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
2217 parquet_writer
2218 .write(&batch)
2219 .map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
2220 }
2221 parquet_writer
2222 .close()
2223 .map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
2224 Ok(())
2225}
2226
2227fn concat_dataset_files(
2236 output_format: &str,
2237 in_paths: &[String],
2238 out_path: &str,
2239) -> Result<(), String> {
2240 if in_paths.is_empty() {
2241 return Err("concat_dataset_files: no chunk files to concatenate".into());
2242 }
2243 if output_format == "parquet" {
2244 concat_parquet_files(in_paths, out_path)
2245 } else {
2246 concat_ipc_files(in_paths, out_path)
2247 }
2248}
2249
2250fn concat_ipc_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
2251 use arrow::ipc::reader::FileReader;
2252
2253 let first_file =
2254 std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
2255 let schema = FileReader::try_new(first_file, None)
2256 .map_err(|e| format!("ipc reader {}: {e}", in_paths[0]))?
2257 .schema();
2258
2259 let out_file =
2260 std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
2261 let mut writer = arrow::ipc::writer::FileWriter::try_new(out_file, &schema)
2262 .map_err(|e| format!("ipc writer {out_path}: {e}"))?;
2263
2264 for path in in_paths {
2265 let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
2266 let reader =
2267 FileReader::try_new(file, None).map_err(|e| format!("ipc reader {path}: {e}"))?;
2268 for batch_result in reader {
2269 let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
2270 writer
2271 .write(&batch)
2272 .map_err(|e| format!("write {out_path}: {e}"))?;
2273 }
2274 }
2275 writer
2276 .finish()
2277 .map_err(|e| format!("finish {out_path}: {e}"))
2278}
2279
2280fn concat_parquet_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
2281 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
2282
2283 let first_file =
2284 std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
2285 let first_builder = ParquetRecordBatchReaderBuilder::try_new(first_file)
2286 .map_err(|e| format!("parquet reader {}: {e}", in_paths[0]))?;
2287 let schema = first_builder.schema().clone();
2288 let meta_kv: Vec<parquet::file::metadata::KeyValue> = first_builder
2289 .metadata()
2290 .file_metadata()
2291 .key_value_metadata()
2292 .cloned()
2293 .unwrap_or_default();
2294
2295 let out_file =
2296 std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
2297 let zstd = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
2298 let props = parquet::file::properties::WriterProperties::builder()
2299 .set_compression(parquet::basic::Compression::ZSTD(zstd))
2300 .set_max_row_group_row_count(Some(1_000_000))
2301 .set_key_value_metadata(Some(meta_kv))
2302 .build();
2303 let mut writer = parquet::arrow::ArrowWriter::try_new(out_file, schema, Some(props))
2304 .map_err(|e| format!("arrow writer {out_path}: {e}"))?;
2305
2306 for path in in_paths {
2307 let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
2308 let reader = ParquetRecordBatchReaderBuilder::try_new(file)
2309 .map_err(|e| format!("parquet reader {path}: {e}"))?
2310 .build()
2311 .map_err(|e| format!("build parquet reader {path}: {e}"))?;
2312 for batch_result in reader {
2313 let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
2314 writer
2315 .write(&batch)
2316 .map_err(|e| format!("write {out_path}: {e}"))?;
2317 }
2318 }
2319 writer
2320 .close()
2321 .map(|_| ())
2322 .map_err(|e| format!("close {out_path}: {e}"))
2323}
2324
2325#[allow(clippy::too_many_arguments)]
2339pub fn run_chunked(
2340 ctx: &Context,
2341 domain: &str,
2342 total_size: usize,
2343 chunk_size: usize,
2344 seed: u64,
2345 difficulty: &str,
2346 hard_neg_ratio: f64,
2347 singleton_master_fraction: f64,
2348 schema: &crate::schema::DomainSchema,
2349 final_run_id: &str,
2350 output_format: &str,
2351 graph_enabled: bool,
2352 graph_format: &str,
2353 only_entity: Option<&str>,
2354 output_dir: &str,
2355 mut progress: Option<&mut dyn FnMut(usize, usize)>,
2356) -> Result<PipelineOutput, String> {
2357 if chunk_size == 0 || chunk_size >= total_size {
2358 let config = crate::schema::build_pipeline_config(
2359 domain,
2360 total_size,
2361 seed,
2362 difficulty,
2363 hard_neg_ratio,
2364 singleton_master_fraction,
2365 schema,
2366 final_run_id,
2367 output_format,
2368 graph_enabled,
2369 graph_format,
2370 only_entity,
2371 )?;
2372 return run_pipeline_with_progress(ctx, &config, output_dir, progress);
2373 }
2374
2375 let n_chunks = total_size.div_ceil(chunk_size);
2376 let tmp_dir = format!("{output_dir}/.dupehell_chunks_{final_run_id}");
2377 std::fs::create_dir_all(&tmp_dir).map_err(|e| format!("create {tmp_dir:?}: {e}"))?;
2378
2379 let mut offsets = ChunkOffsets::default();
2380 let mut dataset_chunks: Vec<String> = Vec::new();
2381 let mut gt_chunks: Vec<String> = Vec::new();
2382 let mut node_chunks: Vec<String> = Vec::new();
2383 let mut edge_chunks: Vec<String> = Vec::new();
2384 let mut agg = PipelineStats {
2385 total_records: 0,
2386 exact_dups: 0,
2387 fuzzy_dups: 0,
2388 hard_negs: 0,
2389 uniques: 0,
2390 masters: 0,
2391 };
2392
2393 for chunk_idx in 0..n_chunks {
2394 let this_size = chunk_size.min(total_size - chunk_idx * chunk_size);
2395 let chunk_seed = seed.wrapping_add(chunk_idx as u64);
2396 let chunk_run_id = format!("{final_run_id}__chunk{chunk_idx:04}");
2397 let chunk_config = crate::schema::build_pipeline_config(
2398 domain,
2399 this_size,
2400 chunk_seed,
2401 difficulty,
2402 hard_neg_ratio,
2403 singleton_master_fraction,
2404 schema,
2405 &chunk_run_id,
2406 output_format,
2407 graph_enabled,
2408 graph_format,
2409 only_entity,
2410 )?;
2411
2412 let already_written = agg.total_records;
2413 let mut chunk_cb = |done: usize, _total: usize| {
2414 if let Some(cb) = &mut progress {
2415 cb(already_written + done, total_size);
2416 }
2417 };
2418 let (out, next_offsets) = run_pipeline_chunked(
2419 ctx,
2420 &chunk_config,
2421 &tmp_dir,
2422 Some(&mut chunk_cb),
2423 Some(&offsets),
2424 )?;
2425 offsets = next_offsets;
2426
2427 agg.total_records = out.stats.total_records;
2432 agg.exact_dups += out.stats.exact_dups;
2433 agg.fuzzy_dups += out.stats.fuzzy_dups;
2434 agg.hard_negs += out.stats.hard_negs;
2435 agg.uniques += out.stats.uniques;
2436 agg.masters += out.stats.masters;
2437
2438 dataset_chunks.push(out.output_files[0].clone());
2439 gt_chunks.push(out.gt_file.clone());
2440 if let Some(n) = out.nodes {
2441 node_chunks.push(n);
2442 }
2443 if let Some(e) = out.edges {
2444 edge_chunks.push(e);
2445 }
2446 }
2447
2448 let dataset_ext = if output_format == "parquet" {
2449 "parquet"
2450 } else {
2451 "ipc"
2452 };
2453 let final_dataset = format!("{output_dir}/{final_run_id}.{dataset_ext}");
2454 concat_dataset_files(output_format, &dataset_chunks, &final_dataset)?;
2455
2456 let final_gt = format!("{output_dir}/{final_run_id}_ground_truth.{dataset_ext}");
2457 concat_dataset_files(output_format, >_chunks, &final_gt)?;
2458
2459 let graph_ext = if graph_format == "parquet" {
2460 "parquet"
2461 } else {
2462 "ipc"
2463 };
2464 let final_nodes = if !node_chunks.is_empty() {
2465 let p = format!("{output_dir}/{final_run_id}_nodes.{graph_ext}");
2466 concat_dataset_files(graph_format, &node_chunks, &p)?;
2467 Some(p)
2468 } else {
2469 None
2470 };
2471 let final_edges = if !edge_chunks.is_empty() {
2472 let p = format!("{output_dir}/{final_run_id}_edges.{graph_ext}");
2473 concat_dataset_files(graph_format, &edge_chunks, &p)?;
2474 Some(p)
2475 } else {
2476 None
2477 };
2478
2479 std::fs::remove_dir_all(&tmp_dir).ok();
2480
2481 Ok(PipelineOutput {
2482 output_files: vec![final_dataset],
2483 gt_file: final_gt,
2484 stats: agg,
2485 nodes: final_nodes,
2486 edges: final_edges,
2487 })
2488}
2489
2490fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
2493 let ts = std::time::SystemTime::now()
2494 .duration_since(std::time::UNIX_EPOCH)
2495 .map(|d| d.as_secs().to_string())
2496 .unwrap_or_else(|_| "unknown".into());
2497 HashMap::from([
2498 (
2499 "dupehell.generator".into(),
2500 format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
2501 ),
2502 (
2503 "dupehell.provenance".into(),
2504 "dupehell-synthetic-data".into(),
2505 ),
2506 ("dupehell.license".into(), "MIT".into()),
2507 (
2508 "dupehell.purpose".into(),
2509 "Educational Use Only — Record Linkage Benchmarking".into(),
2510 ),
2511 (
2512 "dupehell.url".into(),
2513 "https://github.com/vntoinekaio/DupeHell".into(),
2514 ),
2515 ("dupehell.domain".into(), config.domain.clone()),
2516 ("dupehell.size".into(), config.size.to_string()),
2517 ("dupehell.seed".into(), config.seed.to_string()),
2518 ("dupehell.difficulty".into(), config.difficulty.clone()),
2519 ("dupehell.run_id".into(), config.run_id.clone()),
2520 ("dupehell.timestamp".into(), ts),
2521 (
2522 "dupehell.structural_key_columns".into(),
2523 structural_key_columns_note(config),
2524 ),
2525 ])
2526}
2527
2528fn structural_key_columns_note(config: &PipelineConfig) -> String {
2533 let cols: Vec<&str> = config
2534 .entity_plans
2535 .iter()
2536 .filter_map(|p| p.identifier_col.as_deref())
2537 .collect();
2538 format!(
2539 "{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
2540 cols.join(", ")
2541 )
2542}
2543
2544pub(crate) fn low_cardinality_dict_type() -> DataType {
2553 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
2554}
2555
2556pub(crate) struct DictValues {
2573 values: ArrayRef,
2574 index: HashMap<String, i32>,
2575}
2576
2577impl DictValues {
2578 pub(crate) fn new<I, S>(values: I) -> Self
2579 where
2580 I: IntoIterator<Item = S>,
2581 S: Into<String>,
2582 {
2583 let owned: Vec<String> = values.into_iter().map(Into::into).collect();
2584 let index: HashMap<String, i32> = owned
2585 .iter()
2586 .enumerate()
2587 .map(|(i, v)| (v.clone(), i as i32))
2588 .collect();
2589 Self {
2590 values: Arc::new(StringArray::from(owned)),
2591 index,
2592 }
2593 }
2594
2595 fn idx(&self, value: &str) -> i32 {
2596 *self.index.get(value).unwrap_or_else(|| {
2597 panic!("DictValues: value {value:?} was not in the pre-built dictionary")
2598 })
2599 }
2600
2601 fn build(&self, keys: arrow::array::Int32Array) -> ArrayRef {
2602 Arc::new(
2603 arrow::array::DictionaryArray::<arrow::datatypes::Int32Type>::try_new(
2604 keys,
2605 self.values.clone(),
2606 )
2607 .expect("DictValues::build: key array out of range for this dictionary"),
2608 )
2609 }
2610
2611 pub(crate) fn const_array(&self, value: &str, n: usize) -> ArrayRef {
2613 self.build(arrow::array::Int32Array::from(vec![self.idx(value); n]))
2614 }
2615
2616 pub(crate) fn key(&self, value: &str) -> i32 {
2621 self.idx(value)
2622 }
2623
2624 pub(crate) fn finish_keys(&self, keys: arrow::array::Int32Array) -> ArrayRef {
2625 self.build(keys)
2626 }
2627
2628 pub(crate) fn null_array(&self, n: usize) -> ArrayRef {
2639 self.build(arrow::array::Int32Array::new_null(n))
2640 }
2641}
2642
2643fn subtype_dict(config: &PipelineConfig) -> DictValues {
2653 let mut values: Vec<String> = vec!["complete".to_string(), "spanning_tree".to_string()];
2654 for plan in &config.entity_plans {
2655 for remap in &plan.fk_remaps {
2656 values.push(remap.source_col.clone());
2657 }
2658 }
2659 for hn in &config.hard_neg_types {
2660 if let Ok(v) = serde_json::from_str::<serde_json::Value>(&hn.config_json)
2661 && let Some(pattern) = v.get("pattern").and_then(|p| p.as_str())
2662 {
2663 values.push(pattern.to_string());
2664 }
2665 }
2666 values.sort_unstable();
2667 values.dedup();
2668 DictValues::new(values)
2669}
2670
2671fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
2672 let mut field_map: Vec<(String, DataType, bool)> = vec![
2673 ("record_id".to_string(), DataType::Utf8, false),
2674 ("domain".to_string(), low_cardinality_dict_type(), false),
2675 (
2676 "entity_type".to_string(),
2677 low_cardinality_dict_type(),
2678 false,
2679 ),
2680 ("master_id".to_string(), DataType::Utf8, false),
2681 ];
2682 for plan in &config.entity_plans {
2683 if plan.pool_only {
2688 continue;
2689 }
2690 let cols: Vec<serde_json::Value> =
2691 serde_json::from_str(&plan.columns_json).unwrap_or_default();
2692 for col in &cols {
2693 let name = col["name"].as_str().unwrap_or("").to_string();
2694 if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
2695 continue;
2696 }
2697 let is_dict = crate::fast_template::resolve_dict_encoded_column(&name).is_some();
2706 let col_type = if is_dict {
2707 low_cardinality_dict_type()
2708 } else {
2709 col_type_from_request(col)
2710 };
2711 field_map.push((name, col_type, true));
2712 }
2713 }
2714 let fields: Vec<Field> = field_map
2715 .into_iter()
2716 .map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
2717 .collect();
2718 Schema::new(fields).with_metadata(metadata.clone())
2719}
2720
2721fn col_type_from_request(v: &serde_json::Value) -> DataType {
2723 match v["type"].as_str() {
2724 Some("int") => DataType::Int64,
2725 Some("float") => DataType::Float64,
2726 Some("boolean") => DataType::Boolean,
2727 Some("date") | Some("datetime") => DataType::Utf8,
2728 _ => DataType::Utf8,
2729 }
2730}
2731
2732#[allow(clippy::too_many_arguments)]
2735pub(crate) fn add_metadata_and_align(
2736 rb: &RecordBatch,
2737 domain: &str,
2738 entity_type: &str,
2739 domain_dict: &DictValues,
2740 entity_type_dict: &DictValues,
2741 record_ids: &ArrayRef,
2742 master_ids: &ArrayRef,
2743 full_arc: &Arc<Schema>,
2744 col_lookup: &[Option<usize>],
2745 null_cache: &mut FxHashMap<(DataType, usize), ArrayRef>,
2746 const_arr_cache: &mut FxHashMap<(String, usize), ArrayRef>,
2747) -> Result<RecordBatch, String> {
2748 let n = rb.num_rows();
2749 let domain_arr = const_arr_cache
2759 .entry((domain.to_string(), n))
2760 .or_insert_with(|| domain_dict.const_array(domain, n))
2761 .clone();
2762 let et_arr = const_arr_cache
2763 .entry((entity_type.to_string(), n))
2764 .or_insert_with(|| entity_type_dict.const_array(entity_type, n))
2765 .clone();
2766 let rid_arr = record_ids.clone();
2767 let mid_arr = master_ids.clone();
2768
2769 let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
2770 all_arrays.push(rid_arr);
2771 all_arrays.push(domain_arr);
2772 all_arrays.push(et_arr);
2773 all_arrays.push(mid_arr);
2774
2775 for (i, maybe_idx) in col_lookup.iter().enumerate() {
2778 match maybe_idx {
2779 Some(idx) => all_arrays.push(rb.column(*idx).clone()),
2780 None => {
2781 let field = full_arc.field(i + 4);
2782 if let Some(arr) = crate::fast_template::dict_encoded_null_array(field.name(), n) {
2792 all_arrays.push(arr);
2793 continue;
2794 }
2795 let dt = field.data_type();
2796 let arr = null_cache
2797 .entry((dt.clone(), n))
2798 .or_insert_with(|| arrow::array::new_null_array(dt, n));
2799 all_arrays.push(Arc::clone(arr));
2800 }
2801 }
2802 }
2803
2804 RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
2806 .map_err(|e| format!("add_metadata_and_align RecordBatch: {e}"))
2807}
2808
2809fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
2810 use arrow::compute::take;
2811 let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
2812 .map(|i| take(rb.column(i), indices, None))
2813 .collect::<Result<Vec<_>, _>>()
2814 .map_err(|e| format!("take error: {e}"))?;
2815
2816 RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
2817}
2818
2819#[cfg(test)]
2820mod tests {
2821 use super::*;
2822 use crate::schema::{build_pipeline_config, load_schema};
2823
2824 #[test]
2853 #[ignore = "hunt2808 diagnostic scan, not a pass/fail regression test -- run explicitly with --ignored --nocapture"]
2854 fn test_hunt2808_noise_coverage_scan() {
2855 const HELL_NOISE_TYPES: &[&str] = &[
2856 "typo_aggressive",
2857 "visual",
2858 "unicode_pollution",
2859 "names",
2860 "dates",
2861 "identifiers",
2862 "addresses",
2863 "companies",
2864 "extra",
2865 ];
2866
2867 let mut zero_coverage: Vec<(String, String, usize)> = Vec::new();
2868 let mut total_entities = 0usize;
2869 let mut schema_files: Vec<_> = std::fs::read_dir("schemas")
2870 .expect("read schemas/ dir")
2871 .filter_map(|e| e.ok())
2872 .map(|e| e.path())
2873 .filter(|p| p.extension().is_some_and(|e| e == "json"))
2874 .collect();
2875 schema_files.sort();
2876
2877 for path in &schema_files {
2878 let domain = path.file_stem().unwrap().to_string_lossy().to_string();
2879 let text = std::fs::read_to_string(path).unwrap_or_else(|e| panic!("{path:?}: {e}"));
2880 let json: serde_json::Value =
2881 serde_json::from_str(&text).unwrap_or_else(|e| panic!("{path:?}: {e}"));
2882 let Some(entities) = json.get("entities").and_then(|v| v.as_array()) else {
2883 continue;
2884 };
2885 for entity in entities {
2886 total_entities += 1;
2887 let ename = entity
2888 .get("name")
2889 .and_then(|v| v.as_str())
2890 .unwrap_or("?")
2891 .to_string();
2892 let fk_exclude: std::collections::HashSet<&str> = entity
2893 .get("fk_remaps")
2894 .and_then(|v| v.as_array())
2895 .map(|arr| {
2896 arr.iter()
2897 .filter_map(|r| r.get("source_col").and_then(|v| v.as_str()))
2898 .collect()
2899 })
2900 .unwrap_or_default();
2901 let cols = entity
2902 .get("columns")
2903 .and_then(|v| v.as_array())
2904 .cloned()
2905 .unwrap_or_default();
2906
2907 let mut n_matched_categories = 0usize;
2908 for &ntype in HELL_NOISE_TYPES {
2909 let has_match = cols.iter().any(|c| {
2910 let name = c.get("name").and_then(|v| v.as_str()).unwrap_or("");
2911 let ctype = c.get("type").and_then(|v| v.as_str()).unwrap_or("string");
2912 let is_utf8 = !matches!(ctype, "int" | "float" | "boolean");
2913 is_utf8
2914 && !fk_exclude.contains(name)
2915 && noise_type_targets_column(ntype, name)
2916 });
2917 if has_match {
2918 n_matched_categories += 1;
2919 }
2920 }
2921 if n_matched_categories == 0 {
2922 zero_coverage.push((domain.clone(), ename, cols.len()));
2923 }
2924 }
2925 }
2926
2927 eprintln!(
2928 "\n[hunt2808_coverage] {total_entities} entities scanned across {} domains",
2929 schema_files.len()
2930 );
2931 eprintln!(
2932 "[hunt2808_coverage] {} entities with ZERO matching hell noise category:",
2933 zero_coverage.len()
2934 );
2935 for (domain, entity, ncols) in &zero_coverage {
2936 eprintln!("[hunt2808_coverage] {domain}::{entity} ({ncols} columns)");
2937 }
2938 }
2939
2940 #[test]
2947 fn test_no_person_company_name_cross_contamination() {
2948 for col in [
2949 "given_name",
2950 "family_name",
2951 "middle_name",
2952 "first_name",
2953 "last_name",
2954 ] {
2955 assert!(
2956 noise_type_targets_column("names", col),
2957 "{col} should still match 'names'"
2958 );
2959 assert!(
2960 !noise_type_targets_column("companies", col),
2961 "{col} should NOT match 'companies' (person name, not a company name)"
2962 );
2963 }
2964 for col in ["company_name", "legal_name", "trading_name"] {
2965 assert!(
2966 noise_type_targets_column("companies", col),
2967 "{col} should still match 'companies'"
2968 );
2969 assert!(
2970 !noise_type_targets_column("names", col),
2971 "{col} should NOT match 'names' (company name, not a person name)"
2972 );
2973 }
2974 assert!(
2975 !noise_type_targets_column("addresses", "email_address"),
2976 "email_address should NOT match 'addresses' (it's an email, not a postal address)"
2977 );
2978 assert!(
2979 noise_type_targets_column("identifiers", "email_address"),
2980 "email_address should still match 'identifiers' (corrupt_email)"
2981 );
2982 assert!(
2983 noise_type_targets_column("addresses", "residential_address_line1"),
2984 "a real postal address column should still match 'addresses'"
2985 );
2986 }
2987
2988 fn manifest_dir() -> std::path::PathBuf {
2989 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).to_path_buf()
2990 }
2991
2992 fn scratch_dir(name: &str) -> std::path::PathBuf {
2995 let dir = manifest_dir()
2996 .join("target")
2997 .join("pipeline_test_tmp")
2998 .join(format!("{name}_{}", std::process::id()));
2999 let _ = std::fs::remove_dir_all(&dir);
3000 std::fs::create_dir_all(&dir).expect("create scratch dir");
3001 dir
3002 }
3003
3004 fn run_kyc(
3005 output_dir: &std::path::Path,
3006 output_format: &str,
3007 graph_enabled: bool,
3008 ) -> PipelineOutput {
3009 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
3010 let mut ctx = Context::new(
3011 "kyc",
3012 "en",
3013 &manifest_dir().join("assets/pools").to_string_lossy(),
3014 )
3015 .expect("load context");
3016 let config = build_pipeline_config(
3017 "kyc",
3018 200,
3019 42,
3020 "medium",
3021 0.1,
3022 0.3,
3023 &schema,
3024 "pipeline_test_run",
3025 output_format,
3026 graph_enabled,
3027 "parquet",
3028 None,
3029 )
3030 .expect("build config");
3031 ctx.enable_watermark(&config.domain, config.size, config.seed);
3032 run_pipeline(&ctx, &config, &output_dir.to_string_lossy()).expect("run_pipeline")
3033 }
3034
3035 #[test]
3036 fn test_run_pipeline_produces_dataset_and_gt() {
3037 let dir = scratch_dir("basic");
3038 let output = run_kyc(&dir, "parquet", false);
3039
3040 assert_eq!(output.output_files.len(), 1);
3041 for f in &output.output_files {
3042 assert!(std::path::Path::new(f).exists(), "missing dataset file {f}");
3043 }
3044 assert!(std::path::Path::new(&output.gt_file).exists());
3045 assert!(output.nodes.is_none());
3046 assert!(output.edges.is_none());
3047
3048 let s = &output.stats;
3053 assert!(s.total_records >= s.exact_dups + s.fuzzy_dups + s.hard_negs + s.uniques);
3054 assert!(s.total_records > 0);
3055 assert!(s.masters > 0);
3056
3057 let _ = std::fs::remove_dir_all(&dir);
3058 }
3059
3060 #[test]
3061 fn test_run_pipeline_ipc_format() {
3062 let dir = scratch_dir("ipc");
3063 let output = run_kyc(&dir, "ipc", false);
3064 assert!(output.output_files[0].ends_with(".ipc"));
3065 assert!(std::path::Path::new(&output.output_files[0]).exists());
3066 let _ = std::fs::remove_dir_all(&dir);
3067 }
3068
3069 #[test]
3070 fn test_run_pipeline_graph_output() {
3071 let dir = scratch_dir("graph");
3072 let output = run_kyc(&dir, "parquet", true);
3073 let nodes = output.nodes.expect("nodes path");
3074 let edges = output.edges.expect("edges path");
3075 assert!(std::path::Path::new(&nodes).exists());
3076 assert!(std::path::Path::new(&edges).exists());
3077 let _ = std::fs::remove_dir_all(&dir);
3078 }
3079
3080 #[test]
3081 fn test_fk_edge_weight_from_schema_and_default() {
3082 let dir = scratch_dir("fk_weight");
3089 let schema =
3090 load_schema("aviation", &manifest_dir().join("schemas")).expect("load aviation.json");
3091 let mut ctx = Context::new(
3092 "aviation",
3093 "en",
3094 &manifest_dir().join("assets/pools").to_string_lossy(),
3095 )
3096 .expect("load context");
3097 let config = build_pipeline_config(
3098 "aviation",
3099 300,
3100 42,
3101 "medium",
3102 0.1,
3103 0.3,
3104 &schema,
3105 "pipeline_test_fk_weight",
3106 "ipc",
3107 true,
3108 "ipc",
3109 None,
3110 )
3111 .expect("build config");
3112 ctx.enable_watermark(&config.domain, config.size, config.seed);
3113 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
3114
3115 let edges_path = output.edges.expect("edges path");
3116 let file = std::fs::File::open(&edges_path).expect("open edges file");
3117 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
3118
3119 let mut subtypes_seen: std::collections::HashMap<String, f64> =
3120 std::collections::HashMap::new();
3121 for batch in reader {
3122 let batch = batch.expect("read edges batch");
3123 let edge_type_col = arrow::compute::cast(
3127 batch.column_by_name("edge_type").expect("edge_type col"),
3128 &DataType::Utf8,
3129 )
3130 .expect("cast edge_type to Utf8");
3131 let edge_type = edge_type_col.as_string::<i32>();
3132 let subtype_col = arrow::compute::cast(
3133 batch.column_by_name("subtype").expect("subtype col"),
3134 &DataType::Utf8,
3135 )
3136 .expect("cast subtype to Utf8");
3137 let subtype = subtype_col.as_string::<i32>();
3138 let weight = batch
3139 .column_by_name("weight")
3140 .expect("weight col")
3141 .as_any()
3142 .downcast_ref::<arrow::array::Float64Array>()
3143 .expect("weight as f64");
3144 for i in 0..batch.num_rows() {
3145 if edge_type.value(i) == "fk" {
3146 subtypes_seen.insert(subtype.value(i).to_string(), weight.value(i));
3147 }
3148 }
3149 }
3150
3151 assert_eq!(
3152 subtypes_seen.get("aircraft_id"),
3153 Some(&0.5),
3154 "flight->aircraft fk_weight from schema not propagated to edge weight"
3155 );
3156 assert_eq!(
3157 subtypes_seen.get("airline_id"),
3158 Some(&1.0),
3159 "flight->airline fk_weight (schema default) not propagated to edge weight"
3160 );
3161
3162 let _ = std::fs::remove_dir_all(&dir);
3163 }
3164
3165 #[test]
3172 fn test_run_pipeline_dup_cluster_edges_well_formed() {
3173 let dir = scratch_dir("dup_cluster_edges");
3174 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
3179 let mut ctx = Context::new(
3180 "kyc",
3181 "en",
3182 &manifest_dir().join("assets/pools").to_string_lossy(),
3183 )
3184 .expect("load context");
3185 let config = build_pipeline_config(
3186 "kyc",
3187 200,
3188 42,
3189 "medium",
3190 0.1,
3191 0.3,
3192 &schema,
3193 "pipeline_test_dup_cluster_edges",
3194 "ipc",
3195 true,
3196 "ipc",
3197 None,
3198 )
3199 .expect("build config");
3200 ctx.enable_watermark(&config.domain, config.size, config.seed);
3201 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
3202
3203 let edges_path = output.edges.expect("edges path");
3204 let file = std::fs::File::open(&edges_path).expect("open edges file");
3205 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
3206
3207 let mut n_exact = 0usize;
3208 let mut n_fuzzy = 0usize;
3209 for batch in reader {
3210 let batch = batch.expect("read edges batch");
3211 let src = batch
3212 .column_by_name("source_node_id")
3213 .expect("source_node_id col")
3214 .as_any()
3215 .downcast_ref::<arrow::array::StringArray>()
3216 .expect("source_node_id as string");
3217 let tgt = batch
3218 .column_by_name("target_node_id")
3219 .expect("target_node_id col")
3220 .as_any()
3221 .downcast_ref::<arrow::array::StringArray>()
3222 .expect("target_node_id as string");
3223 let edge_type_col = arrow::compute::cast(
3226 batch.column_by_name("edge_type").expect("edge_type col"),
3227 &DataType::Utf8,
3228 )
3229 .expect("cast edge_type to Utf8");
3230 let edge_type = edge_type_col.as_string::<i32>();
3231 let weight = batch
3232 .column_by_name("weight")
3233 .expect("weight col")
3234 .as_any()
3235 .downcast_ref::<arrow::array::Float64Array>()
3236 .expect("weight as f64");
3237
3238 for i in 0..batch.num_rows() {
3239 match edge_type.value(i) {
3240 "exact_dup" => n_exact += 1,
3241 "fuzzy_dup" => n_fuzzy += 1,
3242 _ => continue,
3243 }
3244 assert!(
3249 parse_record_idx(src.value(i)).is_some(),
3250 "malformed source_node_id {:?} on a dup-cluster edge",
3251 src.value(i)
3252 );
3253 assert!(
3254 parse_record_idx(tgt.value(i)).is_some(),
3255 "malformed target_node_id {:?} on a dup-cluster edge",
3256 tgt.value(i)
3257 );
3258 assert!((weight.value(i) - 1.0).abs() < 1e-9);
3259 }
3260 }
3261
3262 assert!(n_exact + n_fuzzy > 0, "no dup-cluster edges emitted");
3267 assert!(output.stats.exact_dups + output.stats.fuzzy_dups > 0);
3268
3269 let _ = std::fs::remove_dir_all(&dir);
3270 }
3271
3272 #[test]
3273 fn test_run_pipeline_deterministic_stats() {
3274 let dir_a = scratch_dir("det_a");
3275 let dir_b = scratch_dir("det_b");
3276 let a = run_kyc(&dir_a, "parquet", false);
3277 let b = run_kyc(&dir_b, "parquet", false);
3278 assert_eq!(a.stats.total_records, b.stats.total_records);
3279 assert_eq!(a.stats.exact_dups, b.stats.exact_dups);
3280 assert_eq!(a.stats.fuzzy_dups, b.stats.fuzzy_dups);
3281 assert_eq!(a.stats.hard_negs, b.stats.hard_negs);
3282 assert_eq!(a.stats.uniques, b.stats.uniques);
3283 assert_eq!(a.stats.masters, b.stats.masters);
3284 let _ = std::fs::remove_dir_all(&dir_a);
3285 let _ = std::fs::remove_dir_all(&dir_b);
3286 }
3287
3288 #[test]
3289 fn test_distribute_by_weight_sums_to_total_and_is_proportional() {
3290 let weights = [500_000usize, 216_666];
3298 let counts = distribute_by_weight(283_334, &weights);
3299 assert_eq!(counts.len(), 2);
3300 assert_eq!(counts.iter().sum::<usize>(), 283_334);
3301 assert!(counts[0] > 0, "first batch got zero duplicates");
3305 assert!(counts[1] > 0, "last batch got zero duplicates");
3306 assert!(counts[0] > counts[1]);
3307
3308 assert_eq!(distribute_by_weight(0, &weights), vec![0, 0]);
3310 assert_eq!(distribute_by_weight(100, &[0, 0]), vec![0, 0]);
3311 }
3312
3313 #[test]
3314 fn test_entity_prefix_unique_per_index() {
3315 let a = entity_prefix(0);
3316 let b = entity_prefix(1);
3317 assert_ne!(a, b);
3318 assert_eq!(a, "E00000");
3319 assert_eq!(entity_prefix(42), "E00042");
3320 }
3321
3322 #[test]
3323 fn test_run_pipeline_hell_multi_pass_touches_multiple_columns() {
3324 let dir = scratch_dir("hell_multipass");
3333 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
3334 let mut ctx = Context::new(
3335 "kyc",
3336 "en",
3337 &manifest_dir().join("assets/pools").to_string_lossy(),
3338 )
3339 .expect("load context");
3340 let config = build_pipeline_config(
3341 "kyc",
3342 5000,
3343 7,
3344 "hell",
3345 0.1,
3346 0.3,
3347 &schema,
3348 "pipeline_test_hell_multipass",
3349 "ipc",
3350 false,
3351 "ipc",
3352 None,
3353 )
3354 .expect("build config");
3355 ctx.enable_watermark(&config.domain, config.size, config.seed);
3356 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
3357
3358 let read_all = |path: &str| -> RecordBatch {
3359 let file = std::fs::File::open(path).unwrap();
3360 let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
3361 let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
3362 let schema = batches[0].schema();
3363 let refs: Vec<Vec<&dyn Array>> = (0..schema.fields().len())
3364 .map(|i| batches.iter().map(|b| b.column(i).as_ref()).collect())
3365 .collect();
3366 let cols: Vec<ArrayRef> = refs
3367 .iter()
3368 .map(|r| arrow::compute::concat(r).unwrap())
3369 .collect();
3370 RecordBatch::try_new(schema, cols).unwrap()
3371 };
3372 let ds = read_all(&output.output_files[0]);
3373 let gt = read_all(&output.gt_file);
3374
3375 let ds_rid = ds.column_by_name("record_id").unwrap().as_string::<i32>();
3376 let ds_mid = ds.column_by_name("master_id").unwrap().as_string::<i32>();
3377 let gt_rid = gt.column_by_name("record_id").unwrap().as_string::<i32>();
3378 let gt_mt_col =
3381 arrow::compute::cast(gt.column_by_name("match_type").unwrap(), &DataType::Utf8)
3382 .unwrap();
3383 let gt_mt = gt_mt_col.as_string::<i32>();
3384
3385 let mut match_type_by_rid: HashMap<&str, &str> = HashMap::new();
3386 for i in 0..gt.num_rows() {
3387 match_type_by_rid.insert(gt_rid.value(i), gt_mt.value(i));
3388 }
3389
3390 let mut master_row_by_mid: HashMap<&str, usize> = HashMap::new();
3397 for i in 0..ds.num_rows() {
3398 master_row_by_mid.entry(ds_mid.value(i)).or_insert(i);
3399 }
3400
3401 let compare_cols: Vec<usize> = (0..ds.num_columns())
3402 .filter(|&i| {
3403 let field = ds.schema().field(i).clone();
3404 matches!(field.data_type(), DataType::Utf8)
3405 && !matches!(
3406 field.name().as_str(),
3407 "record_id" | "domain" | "entity_type" | "master_id"
3408 )
3409 && !field.name().ends_with("_id")
3410 })
3411 .collect();
3412
3413 let mut max_changed_cols = 0usize;
3414 for i in 0..ds.num_rows() {
3415 if match_type_by_rid.get(ds_rid.value(i)) != Some(&"fuzzy_dup") {
3416 continue;
3417 }
3418 let Some(&master_i) = master_row_by_mid.get(ds_mid.value(i)) else {
3419 continue;
3420 };
3421 let mut changed = 0usize;
3422 for &col_i in &compare_cols {
3423 let col = ds.column(col_i).as_string::<i32>();
3424 let a = (!col.is_null(i)).then(|| col.value(i));
3425 let b = (!col.is_null(master_i)).then(|| col.value(master_i));
3426 if a != b {
3427 changed += 1;
3428 }
3429 }
3430 max_changed_cols = max_changed_cols.max(changed);
3431 }
3432
3433 assert!(
3434 max_changed_cols >= 2,
3435 "expected at least one fuzzy_dup row to differ from its master on \
3436 2+ columns (evidence of multiple noise passes), got max {max_changed_cols}"
3437 );
3438
3439 let _ = std::fs::remove_dir_all(&dir);
3440 }
3441
3442 #[test]
3443 fn test_record_id_and_pad_string_roundtrip() {
3444 let rid = record_id_string(7);
3445 assert_eq!(rid, "R-0000000000007");
3446 let pad = pad_string(7);
3447 assert_eq!(pad, "0000000000007");
3448 let strs = record_id_strs(0..3);
3449 assert_eq!(strs.len(), 3);
3450 assert_eq!(strs[0], record_id_string(0));
3451 assert_eq!(strs[2], record_id_string(2));
3452 }
3453
3454 #[test]
3455 fn test_parse_record_idx_roundtrip() {
3456 for i in [0usize, 7, 1_234_567, 300_000_000] {
3457 assert_eq!(parse_record_idx(&record_id_string(i)), Some(i as u64));
3458 }
3459 assert_eq!(parse_record_idx("garbage"), None);
3460 assert_eq!(parse_record_idx("R-1"), None);
3461 assert_eq!(parse_record_idx("X-0000000000007"), None);
3462 }
3463
3464 #[test]
3465 fn test_pack_master_key_roundtrip_and_ordering() {
3466 let mid = |e: usize, i: usize| format!("{}-{}", entity_prefix(e), pad_string(i));
3467 assert!(pack_master_key(&mid(0, 0)).is_some());
3468 assert!(pack_master_key(&mid(0, 1)).unwrap() > pack_master_key(&mid(0, 0)).unwrap());
3470 assert!(
3473 pack_master_key(&mid(1, 0)).unwrap()
3474 > pack_master_key(&mid(0, 9_999_999_999_999)).unwrap()
3475 );
3476 assert_eq!(pack_master_key("garbage"), None);
3477 assert_eq!(pack_master_key("HN-0000000000007"), None);
3478 assert_eq!(pack_master_key("CANARY-000-0-abc"), None);
3479 }
3480
3481 #[test]
3491 #[ignore]
3492 fn hunt_h1_zip_cost_vs_column_count() {
3493 use arrow::array::StringArray;
3494 use arrow::datatypes::{DataType, Field, Schema};
3495
3496 fn make_batch(n_rows: usize, n_cols: usize) -> RecordBatch {
3497 let fields: Vec<Field> = (0..n_cols)
3498 .map(|i| Field::new(format!("col{i}"), DataType::Utf8, false))
3499 .collect();
3500 let cols: Vec<ArrayRef> = (0..n_cols)
3501 .map(|_| {
3502 Arc::new(StringArray::from(
3503 (0..n_rows)
3504 .map(|_| "abcdefghijklmnop".to_string())
3505 .collect::<Vec<_>>(),
3506 )) as ArrayRef
3507 })
3508 .collect();
3509 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
3510 }
3511
3512 const N_ROWS: usize = 100_000;
3513 const N_REPS: usize = 5;
3514 const EXTRA_PASSES: usize = 2;
3518
3519 eprintln!("\n[hunt_h1] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
3520 eprintln!("[hunt_h1] C=total columns, T=1 noised column (fixed) -- H1 predicts cost ∝ C");
3521
3522 for &n_cols in &[8usize, 15, 23] {
3523 let plan_cols = vec!["col0".to_string()];
3524 let all_types: Vec<&str> = vec!["typo"];
3525 let mut total = std::time::Duration::ZERO;
3526
3527 for rep in 0..N_REPS {
3528 let batch = make_batch(N_ROWS, n_cols);
3529 let mut rng = Rng::new(42 + rep as u64);
3530 let t0 = std::time::Instant::now();
3531
3532 let (noisy, mut unchanged, _) =
3533 apply_noise_with_retry(&batch, "typo", &plan_cols, &mut rng, &[]).unwrap();
3534 let mut current = noisy;
3535 for _ in 0..EXTRA_PASSES {
3536 current = apply_extra_pass_per_row(
3537 ¤t,
3538 &mut unchanged,
3539 &all_types,
3540 &mut rng,
3541 &[],
3542 )
3543 .unwrap();
3544 }
3545 std::hint::black_box(¤t);
3546 total += t0.elapsed();
3547 }
3548
3549 let avg = total / N_REPS as u32;
3550 eprintln!(
3551 "[hunt_h1] C={n_cols:2} avg={:>8.3}ms per_column={:>7.4}ms per_row_per_column={:>10.6}us",
3552 avg.as_secs_f64() * 1000.0,
3553 avg.as_secs_f64() * 1000.0 / n_cols as f64,
3554 avg.as_secs_f64() * 1_000_000.0 / (N_ROWS * n_cols) as f64,
3555 );
3556 }
3557 }
3558
3559 #[test]
3576 #[ignore]
3577 fn hunt_h1_residual_kyc_vs_nonprofit() {
3578 use arrow::array::StringArray;
3579 use arrow::datatypes::{DataType, Field, Schema};
3580
3581 fn make_batch(n_rows: usize, col_names: &[&str]) -> RecordBatch {
3582 let fields: Vec<Field> = col_names
3583 .iter()
3584 .map(|n| Field::new(*n, DataType::Utf8, true))
3585 .collect();
3586 let cols: Vec<ArrayRef> = col_names
3587 .iter()
3588 .map(|_| {
3589 Arc::new(StringArray::from(
3590 (0..n_rows)
3591 .map(|_| "abcdefghijklmnop testword".to_string())
3592 .collect::<Vec<_>>(),
3593 )) as ArrayRef
3594 })
3595 .collect();
3596 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
3597 }
3598
3599 const N_ROWS: usize = 100_000;
3600 const N_REPS: usize = 5;
3601 const HELL_TYPES: &[&str] = &[
3602 "typo_aggressive",
3603 "visual",
3604 "unicode_pollution",
3605 "names",
3606 "dates",
3607 "identifiers",
3608 "addresses",
3609 "companies",
3610 "extra",
3611 ];
3612 const HELL_PASSES: usize = 8; let kyc_cols: &[&str] = &[
3615 "party_id",
3616 "given_name",
3617 "family_name",
3618 "middle_name",
3619 "birth_date",
3620 "birth_place",
3621 "nationality",
3622 "email_address",
3623 "mobile_phone",
3624 "personal_administrative_number",
3625 "tax_id",
3626 "residential_address_line1",
3627 "residential_city",
3628 "residential_postal_code",
3629 "residential_state",
3630 "residential_country",
3631 "occupation",
3632 "risk_score",
3633 "document_type",
3634 "document_number",
3635 "source_system",
3636 ];
3637 let nonprofit_cols: &[&str] = &[
3638 "donation_id",
3639 "donor_id",
3640 "campaign_id",
3641 "donation_date",
3642 "payment_method",
3643 "source_system",
3644 ];
3645
3646 eprintln!("\n[hunt_h1_residual] N_ROWS={N_ROWS} N_REPS={N_REPS} HELL_PASSES={HELL_PASSES}");
3647
3648 {
3653 let batch = make_batch(N_ROWS, kyc_cols);
3654 eprintln!("[hunt_h1_residual] per-category breakdown on kyc.natural_person:");
3655 for &ttype in HELL_TYPES {
3656 let mut rng = Rng::new(42);
3657 let t0 = std::time::Instant::now();
3658 let (_, _, tcols) =
3659 apply_noise_with_retry(&batch, ttype, &[], &mut rng, &[]).unwrap();
3660 eprintln!(
3661 "[hunt_h1_residual] {ttype:20} {:>8.3}ms target_cols={:?}",
3662 t0.elapsed().as_secs_f64() * 1000.0,
3663 tcols,
3664 );
3665 }
3666 }
3667
3668 for (label, col_names) in [
3669 ("kyc.natural_person", kyc_cols),
3670 ("nonprofit.donation", nonprofit_cols),
3671 ] {
3672 let mut total = std::time::Duration::ZERO;
3673 let mut total_target_hits = 0usize;
3674
3675 for rep in 0..N_REPS {
3676 let batch = make_batch(N_ROWS, col_names);
3677 let mut rng = Rng::new(42 + rep as u64);
3678 let t0 = std::time::Instant::now();
3679
3680 let (noisy, mut unchanged, t0cols) =
3683 apply_noise_with_retry(&batch, "identifiers", &[], &mut rng, &[]).unwrap();
3684 total_target_hits += t0cols.len();
3685 let mut current = noisy;
3686 for _ in 1..HELL_PASSES {
3687 current = apply_extra_pass_per_row(
3688 ¤t,
3689 &mut unchanged,
3690 HELL_TYPES,
3691 &mut rng,
3692 &[],
3693 )
3694 .unwrap();
3695 }
3696 std::hint::black_box(¤t);
3697 total += t0.elapsed();
3698 }
3699
3700 let avg = total / N_REPS as u32;
3701 eprintln!(
3702 "[hunt_h1_residual] {label:20} C={:2} avg={:>8.3}ms avg_first_pass_target_cols={:.1}",
3703 col_names.len(),
3704 avg.as_secs_f64() * 1000.0,
3705 total_target_hits as f64 / N_REPS as f64,
3706 );
3707 }
3708 }
3709}