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
151const RID_LEN: usize = 15; const PAD_LEN: usize = 13; #[inline]
166pub(crate) fn record_id_string(i: usize) -> String {
167 let mut buf = [0u8; RID_LEN];
168 buf[0] = b'R';
169 buf[1] = b'-';
170 let mut n = i;
171 for j in (2..RID_LEN).rev() {
172 buf[j] = b'0' + (n % 10) as u8;
173 n /= 10;
174 }
175 String::from_utf8(buf.to_vec()).expect("record_id_string: buffer is 'R', '-' and ASCII digits")
176}
177
178#[cfg(test)]
182#[inline]
183pub(crate) fn pad_string(i: usize) -> String {
184 let mut buf = [0u8; PAD_LEN];
185 let mut n = i;
186 for j in (0..PAD_LEN).rev() {
187 buf[j] = b'0' + (n % 10) as u8;
188 n /= 10;
189 }
190 String::from_utf8(buf.to_vec()).expect("pad_string: buffer is always ASCII digits")
191}
192
193#[cfg(test)]
194pub(crate) fn record_id_strs(range: std::ops::Range<usize>) -> Vec<String> {
195 range.map(record_id_string).collect()
196}
197
198#[inline]
204fn write_digits(buf: &mut Vec<u8>, val: u64, width: usize) {
205 let start = buf.len();
206 buf.resize(start + width, b'0');
207 let mut n = val;
208 for j in (0..width).rev() {
209 buf[start + j] = b'0' + (n % 10) as u8;
210 n /= 10;
211 }
212}
213
214pub(crate) fn record_id_array(range: std::ops::Range<usize>) -> ArrayRef {
219 let n = range.len();
220 let mut i = range.start as u64;
221 crate::buf_gen::build_string_array(n, RID_LEN, |buf| {
222 buf.push(b'R');
223 buf.push(b'-');
224 write_digits(buf, i, PAD_LEN);
225 i += 1;
226 })
227}
228
229fn master_id_array_range(prefix: &str, start: usize, n: usize) -> ArrayRef {
233 let width = prefix.len() + 1 + PAD_LEN;
234 let mut i = start as u64;
235 crate::buf_gen::build_string_array(n, width, |buf| {
236 buf.extend_from_slice(prefix.as_bytes());
237 buf.push(b'-');
238 write_digits(buf, i, PAD_LEN);
239 i += 1;
240 })
241}
242
243fn master_id_array_from_indices(
248 prefix: &str,
249 mut indices: impl Iterator<Item = usize>,
250 n: usize,
251) -> ArrayRef {
252 let width = prefix.len() + 1 + PAD_LEN;
253 crate::buf_gen::build_string_array(n, width, |buf| {
254 let idx = indices
255 .next()
256 .expect("master_id_array_from_indices: fewer than n indices");
257 buf.extend_from_slice(prefix.as_bytes());
258 buf.push(b'-');
259 write_digits(buf, idx as u64, PAD_LEN);
260 })
261}
262
263fn hn_master_id_array(start: u64, n: usize) -> ArrayRef {
267 const HN_WIDTH: usize = 12; let mut id = start;
269 crate::buf_gen::build_string_array(n, HN_WIDTH, |buf| {
270 buf.extend_from_slice(b"HN-");
271 write_digits(buf, id, 9);
272 id += 1;
273 })
274}
275
276pub(crate) fn parse_record_idx(rid: &str) -> Option<u64> {
282 if rid.len() != RID_LEN || !rid.starts_with("R-") {
283 return None;
284 }
285 rid[2..].parse().ok()
286}
287
288pub(crate) fn pack_master_key(mid: &str) -> Option<u64> {
302 let bytes = mid.as_bytes();
303 if bytes.len() != 1 + 5 + 1 + PAD_LEN || bytes[0] != b'E' || bytes[6] != b'-' {
304 return None;
305 }
306 let entity_idx: u64 = mid[1..6].parse().ok()?;
307 let local_idx: u64 = mid[7..7 + PAD_LEN].parse().ok()?;
308 Some((entity_idx << 44) | local_idx)
309}
310
311fn distribute_by_weight(total: usize, weights: &[usize]) -> Vec<usize> {
319 let weight_sum: usize = weights.iter().sum();
320 if weight_sum == 0 || total == 0 {
321 return vec![0; weights.len()];
322 }
323 let raw: Vec<f64> = weights
324 .iter()
325 .map(|&w| total as f64 * w as f64 / weight_sum as f64)
326 .collect();
327 let mut counts: Vec<usize> = raw.iter().map(|&r| r as usize).collect();
328 let assigned: usize = counts.iter().sum();
329 let mut remainder = total - assigned;
330 let mut order: Vec<usize> = (0..counts.len()).collect();
331 order.sort_by(|&a, &b| {
332 let fa = raw[a] - counts[a] as f64;
333 let fb = raw[b] - counts[b] as f64;
334 fb.partial_cmp(&fa).unwrap_or(std::cmp::Ordering::Equal)
335 });
336 for &i in &order {
337 if remainder == 0 {
338 break;
339 }
340 counts[i] += 1;
341 remainder -= 1;
342 }
343 counts
344}
345
346pub(crate) fn entity_prefix(entity_idx: usize) -> String {
355 format!("E{:05}", entity_idx)
356}
357
358type FkPoolMap = HashMap<String, RecordBatch>;
361
362struct HnPool {
363 batch: RecordBatch,
364 total_count: usize,
365 record_ids: ArrayRef,
368}
369
370const PERSON_NAME_WORDS: &[&str] = &[
377 "name",
378 "first_name",
386 "last_name",
387 "given",
388 "family",
389 "middle",
390 "operator",
391 "technician",
392 "inspector",
393 "claimant",
394 "consignee",
395 "investigator",
396 "collected_by",
397 "performed_by",
398 "assigned_to",
399];
400
401const COMPANY_NAME_WORDS: &[&str] = &[
407 "company",
408 "legal",
409 "trading",
410 "name",
411 "supplier",
412 "manufacturer",
413 "sponsor",
414 "institution",
415 "journal",
416 "funder",
417 "employer",
418 "law_firm",
419 "affiliation",
420 "reported_by",
421 "buyer",
422 "shipper",
423 "airline",
424];
425
426fn is_specifically_person_name(lower: &str) -> bool {
440 PERSON_NAME_WORDS
441 .iter()
442 .any(|&w| w != "name" && lower.contains(w))
443}
444
445fn is_specifically_company_name(lower: &str) -> bool {
448 COMPANY_NAME_WORDS
449 .iter()
450 .any(|&w| w != "name" && lower.contains(w))
451}
452
453pub(crate) fn noise_type_targets_column(noise_type: &str, col_name: &str) -> bool {
459 let lower = col_name.to_lowercase();
460 match noise_type {
461 "typo" | "typo_aggressive" | "typo_extreme" | "qwerty_azerty" | "visual" | "homoglyph"
462 | "unicode_pollution" | "ocr_errors" | "case_swap" | "char_dropout" | "language_mix"
463 | "blocking_fail" => {
464 if lower.contains("email") || lower.contains("ip_address") {
469 return false;
470 }
471 contains_any(&lower, &["address", "street", "city"])
476 || contains_any(&lower, PERSON_NAME_WORDS)
477 || contains_any(&lower, COMPANY_NAME_WORDS)
478 }
479 "names" | "nickname" | "initials" | "partial" | "name_compound" | "swap" | "full_swap" => {
480 contains_any(&lower, PERSON_NAME_WORDS) && !is_specifically_company_name(&lower)
486 }
487 "dates" | "date_error" | "date_chaotic" | "date_format_mix" | "age_impossible" => {
488 contains_any(&lower, &["date", "birth", "incorporation", "founding"])
489 }
490 "missing" | "missing_pattern" => {
491 contains_any(&lower, &["phone", "email", "mobile", "address", "street"])
492 }
493 "identifiers"
494 | "corrupt_email"
495 | "corrupt_phone"
496 | "corrupt_national_id"
497 | "corrupt_siren"
498 | "national_id_corrupt"
499 | "phone_corrupt"
500 | "email_corrupt"
501 | "siren_corrupt" => contains_any(
502 &lower,
503 &[
504 "email",
505 "phone",
506 "ssn",
507 "pan",
508 "passport",
509 "account",
510 "medicare",
511 "license",
512 "siren",
513 "national_id",
514 "registration",
521 ],
522 ),
523 "extra"
524 | "name_null"
525 | "dob_null"
526 | "blocking_fail_initial"
527 | "blocking_fail_partial"
528 | "fuzzy_match"
529 | "phonetic" => {
530 if lower.contains("email") || lower.contains("ip_address") {
536 return false;
537 }
538 contains_any(&lower, &["address", "street", "city", "note", "comment"])
539 || contains_any(&lower, PERSON_NAME_WORDS)
540 || contains_any(&lower, COMPANY_NAME_WORDS)
541 }
542 "companies" | "acronym" | "legal_form_drop" | "word_dropout" | "company_scramble" => {
543 contains_any(&lower, COMPANY_NAME_WORDS) && !is_specifically_person_name(&lower)
549 }
550 "addresses" | "address_scramble" | "postal_corrupt" => {
551 !lower.contains("ip_address")
567 && !lower.contains("email")
568 && contains_any(
569 &lower,
570 &["address", "street", "postal", "city", "state", "country"],
571 )
572 }
573 "exact" | "english_name" | "estonian_name" | "lithuanian_name" | "slovak_name"
574 | "serbian_name" | "norwegian_name" | "swedish_name" | "dutch_name" | "czech_name"
575 | "albanian_name" | "polish_name" | "romanian_name" | "hungarian_name" | "german_name"
576 | "italian_name" | "spanish_name" | "portuguese_name" | "combo_hard" | "combo_extreme"
577 | "combo_ultimate" | "french_address" => {
578 false
580 }
581 _ => true,
583 }
584}
585
586fn match_noise_columns(schema: &Schema, noise_type: &str, exclude_cols: &[String]) -> Vec<String> {
587 let mut matched: Vec<String> = Vec::new();
588 for field in schema.fields() {
589 if !matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) {
590 continue;
591 }
592 if exclude_cols.contains(field.name()) {
593 continue;
594 }
595 if noise_type_targets_column(noise_type, field.name()) {
596 matched.push(field.name().clone());
597 }
598 }
599 if matches!(noise_type, "swap" | "full_swap") && matched.len() >= 2 {
601 matched.truncate(2);
602 }
603 matched.sort();
604 matched.dedup();
605 matched
606}
607
608fn contains_any(s: &str, patterns: &[&str]) -> bool {
609 patterns.iter().any(|p| s.contains(p))
610}
611
612fn apply_noise_to_batch(
619 rb: &RecordBatch,
620 noise_type: &str,
621 plan_cols: &[String],
622 rng: &mut Rng,
623 exclude_cols: &[String],
624) -> Result<(RecordBatch, Vec<String>), String> {
625 let schema = rb.schema();
626 let n_cols = rb.num_columns();
627
628 let target_cols: Vec<String> = if plan_cols.is_empty() {
630 match_noise_columns(&schema, noise_type, exclude_cols)
631 } else {
632 plan_cols.to_vec()
633 };
634
635 let mut new_columns: Vec<Option<ArrayRef>> = (0..n_cols).map(|_| None).collect();
637
638 let get_col = |new_columns: &[Option<ArrayRef>], idx: usize, rb: &RecordBatch| -> ArrayRef {
639 new_columns[idx]
640 .clone()
641 .unwrap_or_else(|| rb.column(idx).clone())
642 };
643
644 match noise_type {
645 "blocking_fail" => {
646 for col_name in &target_cols {
647 let lower = col_name.to_lowercase();
648 if let Ok(idx) = schema.index_of(col_name) {
649 let col = get_col(&new_columns, idx, rb);
650 if lower.contains("first") || lower.contains("given") {
651 new_columns[idx] =
652 Some(crate::noise::extra::apply_blocking_initial(&*col, rng));
653 } else {
654 new_columns[idx] =
655 Some(crate::noise::extra::apply_blocking_partial(&*col, rng));
656 }
657 }
658 }
659 }
660 "swap" | "full_swap" => {
661 if target_cols.len() >= 2 {
662 let name_a = &target_cols[0];
663 let name_b = &target_cols[1];
664 if let (Ok(idx_a), Ok(idx_b)) = (schema.index_of(name_a), schema.index_of(name_b)) {
665 let a = get_col(&new_columns, idx_a, rb);
666 let b = get_col(&new_columns, idx_b, rb);
667 new_columns[idx_a] = Some(b);
668 new_columns[idx_b] = Some(a);
669 }
670 }
671 }
672 "missing_pattern" => {
673 let mut phone_idx: Option<usize> = None;
674 let mut email_idx: Option<usize> = None;
675 for col_name in &target_cols {
676 let lower = col_name.to_lowercase();
677 if let Ok(idx) = schema.index_of(col_name) {
678 if lower.contains("phone") || lower.contains("mobile") {
679 phone_idx = Some(idx);
680 } else if lower.contains("email") {
681 email_idx = Some(idx);
682 }
683 }
684 }
685 if let (Some(pi), Some(ei)) = (phone_idx, email_idx) {
686 let phone_arr = get_col(&new_columns, pi, rb);
687 let email_arr = get_col(&new_columns, ei, rb);
688 let phone_s = phone_arr.as_string::<i32>();
689 let email_s = email_arr.as_string::<i32>();
690 let n = phone_s.len();
691 let mut builder = arrow::array::StringBuilder::with_capacity(n, 32);
692 for i in 0..n {
693 let pv = if phone_s.is_null(i) {
694 ""
695 } else {
696 phone_s.value(i)
697 };
698 let ev = if email_s.is_null(i) {
699 ""
700 } else {
701 email_s.value(i)
702 };
703 if !ev.is_empty() && rng.next_usize(2) == 0 {
704 builder.append_value("");
705 } else {
706 builder.append_value(pv);
707 }
708 }
709 new_columns[pi] = Some(Arc::new(builder.finish()));
710 }
711 }
712 _ => {
713 for col_name in &target_cols {
714 if let Ok(col_idx) = schema.index_of(col_name) {
715 let col = get_col(&new_columns, col_idx, rb);
716 let result = crate::noise::apply_noise_to_column(&*col, noise_type, rng)
717 .map_err(|e| format!("noise '{noise_type}' on '{col_name}': {e}"))?;
718 new_columns[col_idx] = Some(result);
719 }
720 }
721 }
722 }
723
724 let final_columns: Vec<ArrayRef> = (0..n_cols)
726 .map(|i| {
727 new_columns[i]
728 .clone()
729 .unwrap_or_else(|| rb.column(i).clone())
730 })
731 .collect();
732
733 let batch =
734 RecordBatch::try_new(schema, final_columns).map_err(|e| format!("RecordBatch: {e}"))?;
735 Ok((batch, target_cols))
736}
737
738fn unchanged_row_mask(
749 orig: &RecordBatch,
750 noised: &RecordBatch,
751 target_cols: &[String],
752) -> Result<Vec<bool>, String> {
753 let n = orig.num_rows();
754 let mut changed = vec![false; n];
755 let schema = orig.schema();
756 for col_name in target_cols {
757 let Ok(idx) = schema.index_of(col_name) else {
758 continue;
759 };
760 let a = orig.column(idx).as_string::<i32>();
761 let b = noised.column(idx).as_string::<i32>();
762 for (i, is_changed) in changed.iter_mut().enumerate() {
763 if *is_changed {
764 continue;
765 }
766 let av = (!a.is_null(i)).then(|| a.value(i));
767 let bv = (!b.is_null(i)).then(|| b.value(i));
768 if av != bv {
769 *is_changed = true;
770 }
771 }
772 }
773 for c in changed.iter_mut() {
774 *c = !*c;
775 }
776 Ok(changed)
777}
778
779fn apply_noise_with_retry(
802 orig: &RecordBatch,
803 noise_type: &str,
804 plan_cols: &[String],
805 rng: &mut Rng,
806 exclude_cols: &[String],
807) -> Result<(RecordBatch, Vec<bool>, Vec<String>), String> {
808 const MAX_NOISE_ATTEMPTS: usize = 4;
809 let (mut noisy, target_cols) =
810 apply_noise_to_batch(orig, noise_type, plan_cols, rng, exclude_cols)?;
811 if target_cols.is_empty() {
812 return Ok((noisy, vec![true; orig.num_rows()], target_cols));
813 }
814 let schema = noisy.schema();
815 let target_idxs: Vec<usize> = target_cols
816 .iter()
817 .filter_map(|c| schema.index_of(c).ok())
818 .collect();
819 let mut unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
820 let mut attempt = 1;
821 while unchanged.iter().any(|&u| u) && attempt < MAX_NOISE_ATTEMPTS {
822 let mut retry_rng = Rng::new(rng.next_u64());
823 let (retried, _) =
824 apply_noise_to_batch(orig, noise_type, plan_cols, &mut retry_rng, exclude_cols)?;
825 let mask = arrow::array::BooleanArray::from_iter(unchanged.iter().map(|&u| Some(u)));
826 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
827 for &i in &target_idxs {
828 merged_cols[i] =
829 arrow::compute::kernels::zip::zip(&mask, retried.column(i), noisy.column(i))
830 .map_err(|e| format!("zip noise retry col {i}: {e}"))?;
831 }
832 noisy = RecordBatch::try_new(noisy.schema(), merged_cols)
833 .map_err(|e| format!("rebuild retried dup batch: {e}"))?;
834 unchanged = unchanged_row_mask(orig, &noisy, &target_cols)?;
835 attempt += 1;
836 }
837 Ok((noisy, unchanged, target_cols))
838}
839
840fn apply_extra_pass_per_row(
867 noisy: &RecordBatch,
868 unchanged: &mut [bool],
869 all_types: &[&str],
870 rng: &mut Rng,
871 exclude_cols: &[String],
872) -> Result<RecordBatch, String> {
873 let n = noisy.num_rows();
874 if all_types.is_empty() || n == 0 {
875 return Ok(noisy.clone());
876 }
877 let per_row_type: Vec<usize> = (0..n).map(|_| rng.next_usize(all_types.len())).collect();
878
879 let mut groups: Vec<Vec<u32>> = vec![Vec::new(); all_types.len()];
898 let mut row_group: Vec<u32> = Vec::with_capacity(n);
899 let mut row_pos: Vec<u32> = Vec::with_capacity(n);
900 for (i, &t) in per_row_type.iter().enumerate() {
901 row_pos.push(groups[t].len() as u32);
902 row_group.push(t as u32);
903 groups[t].push(i as u32);
904 }
905
906 let schema = noisy.schema();
907 let mut pass_unchanged = vec![true; n];
908 let mut noised_by_group: Vec<Option<RecordBatch>> = vec![None; all_types.len()];
909 let mut group_target_idxs: Vec<Vec<usize>> = vec![Vec::new(); all_types.len()];
915 let mut touched_cols: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
916
917 for (type_idx, idxs) in groups.into_iter().enumerate() {
918 if idxs.is_empty() {
919 continue;
920 }
921 let ttype = all_types[type_idx];
922 let idx_arr = UInt64Array::from_iter_values(idxs.iter().map(|&i| i as u64));
923 let subset = pick_rows(noisy, &idx_arr)?;
924 let (noised_subset, subset_unchanged, target_cols) =
925 apply_noise_with_retry(&subset, ttype, &[], rng, exclude_cols)?;
926 let target_idxs: Vec<usize> = target_cols
927 .iter()
928 .filter_map(|c| schema.index_of(c).ok())
929 .collect();
930 touched_cols.extend(target_idxs.iter().copied());
931 group_target_idxs[type_idx] = target_idxs;
932 for (local_pos, &orig_i) in idxs.iter().enumerate() {
933 pass_unchanged[orig_i as usize] = subset_unchanged[local_pos];
934 }
935 noised_by_group[type_idx] = Some(noised_subset);
936 }
937
938 let mut merged_cols: Vec<ArrayRef> = noisy.columns().to_vec();
939 for col_idx in touched_cols {
940 let orig_col = noisy.column(col_idx).as_string::<i32>();
941 let mut builder = StringBuilder::with_capacity(n, n * 16);
942 for i in 0..n {
943 let g = row_group[i] as usize;
944 if group_target_idxs[g].contains(&col_idx) {
945 let pos = row_pos[i] as usize;
946 let noised_col = noised_by_group[g]
947 .as_ref()
948 .expect("row's group has at least this row, so it was populated above")
949 .column(col_idx)
950 .as_string::<i32>();
951 if noised_col.is_null(pos) {
952 builder.append_null();
953 } else {
954 builder.append_value(noised_col.value(pos));
955 }
956 } else if orig_col.is_null(i) {
957 builder.append_null();
958 } else {
959 builder.append_value(orig_col.value(i));
960 }
961 }
962 merged_cols[col_idx] = Arc::new(builder.finish());
963 }
964 let current = RecordBatch::try_new(noisy.schema(), merged_cols)
965 .map_err(|e| format!("rebuild extra-pass batch: {e}"))?;
966
967 for (u, pu) in unchanged.iter_mut().zip(pass_unchanged.iter()) {
968 *u = *u && *pu;
969 }
970 Ok(current)
971}
972
973fn topological_sort(
977 plans: &[EntityPlan],
978 name_to_idx: &HashMap<&str, usize>,
979) -> Result<Vec<usize>, String> {
980 let n = plans.len();
981 let mut in_degree = vec![0usize; n];
982 let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
983 for i in 0..n {
984 for remap in &plans[i].fk_remaps {
985 let Some(&target) = name_to_idx.get(remap.target_entity.as_str()) else {
986 continue;
987 };
988 if target == i {
989 continue;
990 }
991 adj[target].push(i); in_degree[i] += 1;
993 }
994 }
995 let mut queue: Vec<usize> = (0..n).filter(|i| in_degree[*i] == 0).collect();
996 let mut order = Vec::with_capacity(n);
997 let mut idx = 0;
998 while idx < queue.len() {
999 let v = queue[idx];
1000 idx += 1;
1001 order.push(v);
1002 for &next in &adj[v] {
1003 in_degree[next] -= 1;
1004 if in_degree[next] == 0 {
1005 queue.push(next);
1006 }
1007 }
1008 }
1009 if order.len() != n {
1010 return Err("FK dependency cycle detected among entities".into());
1011 }
1012 Ok(order)
1013}
1014
1015pub(crate) enum DatasetWriter {
1019 Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
1020 Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
1021}
1022
1023impl DatasetWriter {
1024 fn new(
1025 output_format: &str,
1026 path: &str,
1027 schema: &Arc<Schema>,
1028 metadata: &HashMap<String, String>,
1029 ) -> Result<Self, String> {
1030 let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
1031 if output_format == "parquet" {
1032 use parquet::basic::{Compression, ZstdLevel};
1033 use parquet::file::properties::WriterProperties;
1034 let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
1035 let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
1036 .iter()
1037 .map(|(k, v)| parquet::file::metadata::KeyValue {
1038 key: k.clone(),
1039 value: Some(v.clone()),
1040 })
1041 .collect();
1042 let props = WriterProperties::builder()
1043 .set_compression(Compression::ZSTD(zstd))
1044 .set_max_row_group_row_count(Some(1_000_000))
1048 .set_key_value_metadata(Some(meta_kv))
1049 .build();
1050 let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
1051 .map_err(|e| format!("ArrowWriter error: {e}"))?;
1052 Ok(DatasetWriter::Parquet(Box::new(writer)))
1053 } else {
1054 let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
1055 .map_err(|e| format!("IPC FileWriter error: {e}"))?;
1056 Ok(DatasetWriter::Ipc(Box::new(writer)))
1057 }
1058 }
1059
1060 pub(crate) fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
1061 match self {
1062 DatasetWriter::Ipc(w) => w.write(batch).map_err(|e| format!("write ipc: {e}")),
1063 DatasetWriter::Parquet(w) => w.write(batch).map_err(|e| format!("write parquet: {e}")),
1064 }
1065 }
1066
1067 fn finish(self) -> Result<(), String> {
1068 match self {
1069 DatasetWriter::Ipc(mut w) => w.finish().map_err(|e| format!("finish ipc: {e}")),
1070 DatasetWriter::Parquet(w) => w
1071 .close()
1072 .map(|_| ())
1073 .map_err(|e| format!("close parquet: {e}")),
1074 }
1075 }
1076}
1077
1078fn log_rss(label: &str) {
1090 if !log::log_enabled!(log::Level::Debug) {
1091 return;
1092 }
1093 let pid = sysinfo::Pid::from_u32(std::process::id());
1094 let mut sys = sysinfo::System::new();
1095 sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
1096 if let Some(p) = sys.process(pid) {
1097 log::debug!("[hunt_rss] {label}: rss={:.1} Mo", p.memory() as f64 / 1e6);
1098 }
1099}
1100
1101pub fn run_pipeline(
1104 ctx: &Context,
1105 config: &PipelineConfig,
1106 output_dir: &str,
1107) -> Result<PipelineOutput, String> {
1108 run_pipeline_with_progress(ctx, config, output_dir, None)
1109}
1110
1111pub fn run_pipeline_with_progress(
1118 ctx: &Context,
1119 config: &PipelineConfig,
1120 output_dir: &str,
1121 progress: Option<&mut dyn FnMut(usize, usize)>,
1122) -> Result<PipelineOutput, String> {
1123 run_pipeline_chunked(ctx, config, output_dir, progress, None).map(|(out, _)| out)
1124}
1125
1126pub fn run_pipeline_chunked(
1133 ctx: &Context,
1134 config: &PipelineConfig,
1135 output_dir: &str,
1136 mut progress: Option<&mut dyn FnMut(usize, usize)>,
1137 offsets: Option<&ChunkOffsets>,
1138) -> Result<(PipelineOutput, ChunkOffsets), String> {
1139 let offsets = offsets.cloned().unwrap_or_default();
1140 std::fs::create_dir_all(output_dir)
1141 .map_err(|e| format!("create output directory {output_dir:?}: {e}"))?;
1142 let t_start = std::time::Instant::now();
1143
1144 let t_alloc = t_start.elapsed().as_secs_f64();
1145 log_rss("start");
1146
1147 let t1e = std::time::Instant::now();
1151
1152 let metadata = build_metadata_map(config);
1158 let full_arc = Arc::new(build_full_schema(config, &metadata));
1160 let domain_dict = DictValues::new([config.domain.clone()]);
1168 let entity_type_dict = DictValues::new(
1169 config
1170 .entity_plans
1171 .iter()
1172 .filter(|p| !p.pool_only)
1173 .map(|p| p.name.clone()),
1174 );
1175 let dataset_ext = if config.output_format == "parquet" {
1176 "parquet"
1177 } else {
1178 "ipc"
1179 };
1180 let dataset_path = format!("{}/{}.{}", output_dir, config.run_id, dataset_ext);
1181 let mut writer =
1182 DatasetWriter::new(&config.output_format, &dataset_path, &full_arc, &metadata)?;
1183
1184 let name_to_idx: HashMap<&str, usize> = config
1186 .entity_plans
1187 .iter()
1188 .enumerate()
1189 .map(|(i, p)| (p.name.as_str(), i))
1190 .collect();
1191 let entity_order = topological_sort(&config.entity_plans, &name_to_idx)?;
1192
1193 let mut fk_pools: FkPoolMap = HashMap::new();
1195 let mut hn_pools: HashMap<String, HnPool> = HashMap::new();
1196
1197 let gt_ext = if config.output_format == "parquet" {
1201 "parquet"
1202 } else {
1203 "ipc"
1204 };
1205 let gt_path = format!("{}/{}_ground_truth.{}", output_dir, config.run_id, gt_ext);
1206 let gt_draft_path = format!("{}/{}_gt_draft.ipc", output_dir, config.run_id);
1207 let mut gt_acc = crate::gt::GtAccumulator::new(>_draft_path)?;
1208
1209 let mut global_rid_offset = offsets.rid_offset;
1210
1211 const GRAPH_MAX_CLUSTER_EDGES: usize = 10_000;
1213 let graph_fmt = crate::graph_gen::GraphFormat::from_str(&config.graph_format);
1214 let mut node_writer: Option<crate::graph_gen::NodeWriter> = None;
1215 let mut edge_writer: Option<crate::graph_gen::EdgeWriter> = None;
1216 let mut graph_nodes_path: Option<String> = None;
1217 let mut graph_edges_path: Option<String> = None;
1218 if config.graph_enabled {
1219 let nodes_ipc = format!("{}/{}_nodes.ipc", output_dir, config.run_id);
1220 let edges_ipc = format!("{}/{}_edges.ipc", output_dir, config.run_id);
1221 node_writer = Some(
1222 crate::graph_gen::NodeWriter::new(&nodes_ipc, &full_arc, &metadata)
1223 .map_err(|e| format!("init node writer: {e}"))?,
1224 );
1225 edge_writer = Some(
1226 crate::graph_gen::EdgeWriter::new(&edges_ipc, &metadata, subtype_dict(config))
1227 .map_err(|e| format!("init edge writer: {e}"))?,
1228 );
1229 graph_nodes_path = Some(nodes_ipc);
1230 graph_edges_path = Some(edges_ipc);
1231 }
1232
1233 let mut _t_write = 0.0f64;
1234 let mut _t_meta = 0.0f64;
1235 let mut _t_dup = 0.0f64;
1236 let mut _write_calls = 0u64;
1237 let hn_pool_sizing: std::collections::HashMap<&str, usize> =
1239 config
1240 .hard_neg_types
1241 .iter()
1242 .fold(std::collections::HashMap::new(), |mut m, hn| {
1243 let entry = m.entry(hn.entity_type.as_str()).or_insert(0);
1244 *entry = (*entry).max(hn.count);
1245 m
1246 });
1247 let hn_max_pool: usize = hn_pool_sizing.values().max().copied().unwrap_or(50_000) * 2 + 1_000;
1248 let pool_only_names: std::collections::HashSet<&str> = config
1264 .entity_plans
1265 .iter()
1266 .filter(|p| p.pool_only)
1267 .map(|p| p.name.as_str())
1268 .collect();
1269
1270 let fk_targets: std::collections::HashSet<&str> = config
1271 .entity_plans
1272 .iter()
1273 .flat_map(|p| p.fk_remaps.iter().map(|r| r.target_entity.as_str()))
1274 .collect();
1275
1276 for &plan_idx in &entity_order {
1277 let plan = &config.entity_plans[plan_idx];
1278 if plan.n_base == 0 {
1279 continue;
1280 }
1281
1282 let prefix = entity_prefix(plan_idx);
1283 let master_base = offsets.master_base(plan_idx);
1284 let parsed_columns = crate::entity_gen::parse_columns(&plan.columns_json)?;
1287
1288 let fk_targeted = fk_targets.contains(plan.name.as_str());
1290 let mut fk_builder = arrow::array::StringBuilder::new();
1291 let mut fk_rid_builder = arrow::array::StringBuilder::new();
1292 let mut fk_count: usize = 0;
1293 let hn_targeted = hn_pool_sizing.contains_key(plan.name.as_str());
1298 let mut hn_slices: Vec<RecordBatch> = Vec::new();
1299 let mut hn_slice_rids: Vec<ArrayRef> = Vec::new();
1300 let mut batch_rng = Rng::new(config.seed.wrapping_add(100));
1301 let mut fk_rng = Rng::new(config.seed.wrapping_add(42));
1302
1303 let batch_bounds: Vec<(usize, usize)> = (0..plan.n_base)
1309 .step_by(crate::entity_gen::BATCH_SIZE)
1310 .map(|offset| {
1311 (
1312 offset,
1313 (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE),
1314 )
1315 })
1316 .collect();
1317 let batch_weights: Vec<usize> = batch_bounds.iter().map(|&(_, bn)| bn).collect();
1318 let per_batch_noise_counts: Vec<Vec<usize>> = plan
1319 .noise_types
1320 .iter()
1321 .map(|n| distribute_by_weight(n.count, &batch_weights))
1322 .collect();
1323
1324 let mut col_lookup: Option<Vec<Option<usize>>> = None;
1326 let mut null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1328 let mut const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1329
1330 for (batch_idx, offset) in (0..plan.n_base)
1332 .step_by(crate::entity_gen::BATCH_SIZE)
1333 .enumerate()
1334 {
1335 let batch_n = (plan.n_base - offset).min(crate::entity_gen::BATCH_SIZE);
1336 let batch_seed = config.seed.wrapping_add(batch_idx as u64 * 1000 + 1000);
1337
1338 let rb = crate::entity_gen::generate_entity_batch_parsed(
1340 ctx,
1341 &parsed_columns,
1342 batch_n,
1343 batch_seed,
1344 offset,
1345 )?;
1346
1347 let mut fk_edges_by_remap: Vec<(String, ArrayRef, f64)> = Vec::new();
1349 let rb = if !plan.fk_remaps.is_empty() {
1350 let mut r = rb;
1351 for remap in &plan.fk_remaps {
1352 if let Some(pool) = fk_pools.get(&remap.target_entity) {
1353 let want_edge = config.graph_enabled
1357 && !pool_only_names.contains(remap.target_entity.as_str());
1358 let (remapped, target_rids) = crate::fk_remap::fk_remap_batch(
1359 &r,
1360 pool,
1361 &remap.source_col,
1362 &mut fk_rng,
1363 want_edge,
1364 )?;
1365 r = remapped;
1366 if want_edge && let Some(rids) = target_rids {
1367 fk_edges_by_remap.push((
1368 remap.source_col.clone(),
1369 rids,
1370 remap.fk_weight,
1371 ));
1372 }
1373 }
1374 }
1375 r
1376 } else {
1377 rb
1378 };
1379
1380 if fk_targeted
1383 && fk_count < FK_POOL_CAP
1384 && let Some(ref id_col) = plan.identifier_col
1385 && let Some(col) = rb.column_by_name(id_col)
1386 {
1387 let s = col.as_string::<i32>();
1388 for i in 0..s.len() {
1389 if fk_count >= FK_POOL_CAP {
1390 break;
1391 }
1392 if !s.is_null(i) {
1393 fk_builder.append_value(s.value(i));
1394 if config.graph_enabled {
1395 fk_rid_builder.append_value(record_id_string(global_rid_offset + i));
1396 }
1397 fk_count += 1;
1398 }
1399 }
1400 }
1401
1402 if plan.pool_only {
1408 continue;
1409 }
1410
1411 let hn_accum: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1413 if hn_targeted && hn_accum < hn_max_pool {
1414 let take = (hn_max_pool - hn_accum).min(batch_n);
1415 hn_slices.push(rb.slice(0, take));
1416 if config.graph_enabled {
1420 hn_slice_rids
1421 .push(record_id_array(global_rid_offset..global_rid_offset + take));
1422 }
1423 }
1424
1425 if col_lookup.is_none() {
1427 let plan_schema = rb.schema();
1428 col_lookup = Some(
1429 full_arc
1430 .fields()
1431 .iter()
1432 .skip(4)
1433 .map(|f| plan_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1434 .collect(),
1435 );
1436 }
1437
1438 let batch_mid_arr = master_id_array_range(&prefix, master_base + offset, batch_n);
1440
1441 let rid_arr = record_id_array(global_rid_offset..global_rid_offset + batch_n);
1443 let t_m0 = std::time::Instant::now();
1444 let base_rb = add_metadata_and_align(
1445 &rb,
1446 &config.domain,
1447 &plan.name,
1448 &domain_dict,
1449 &entity_type_dict,
1450 &rid_arr,
1451 &batch_mid_arr,
1452 &full_arc,
1453 col_lookup.as_ref().unwrap(),
1454 &mut null_cache,
1455 &mut const_arr_cache,
1456 )?;
1457 _t_meta += t_m0.elapsed().as_secs_f64();
1458
1459 if config.graph_enabled
1461 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1462 {
1463 nw.write_batch(&base_rb)
1464 .map_err(|e| format!("write node: {e}"))?;
1465 let fk_edges_by_remap: Vec<_> = fk_edges_by_remap
1466 .iter()
1467 .map(|(subtype, target_rids, weight)| {
1468 (subtype, target_rids.as_string::<i32>(), *weight)
1469 })
1470 .collect();
1471 let rid_str_arr = rid_arr.as_string::<i32>();
1472 for i in 0..batch_n {
1473 let rid = rid_str_arr.value(i);
1474 for (subtype, target_rids, weight) in &fk_edges_by_remap {
1475 ew.push(rid, target_rids.value(i), "fk", subtype, *weight)
1476 .map_err(|e| format!("push fk edge: {e}"))?;
1477 }
1478 }
1479 }
1480
1481 let t_w0 = std::time::Instant::now();
1482 writer
1483 .write(&base_rb)
1484 .map_err(|e| format!("write base: {e}"))?;
1485 _t_write += t_w0.elapsed().as_secs_f64();
1486 _write_calls += 1;
1487
1488 gt_acc.push_base_batch(base_rb.column(0), base_rb.column(2), base_rb.column(3))?;
1490
1491 global_rid_offset += batch_n;
1492 if let Some(cb) = &mut progress {
1493 cb(global_rid_offset, config.size);
1494 }
1495
1496 let has_dups: bool = per_batch_noise_counts
1504 .iter()
1505 .any(|counts| counts[batch_idx] > 0);
1506 if has_dups {
1507 let t_d0 = std::time::Instant::now();
1508 let mut dup_batches: Vec<RecordBatch> = Vec::new();
1509 let mut dup_master_idx_buf: Vec<usize> = Vec::new();
1510
1511 let fk_exclude_cols: Vec<String> = plan
1513 .fk_remaps
1514 .iter()
1515 .map(|r| r.source_col.clone())
1516 .collect();
1517
1518 let ndata: Vec<(UInt64Array, u64, &str, &[String], usize)> = plan
1520 .noise_types
1521 .iter()
1522 .enumerate()
1523 .filter_map(|(ni, n)| {
1524 let cnt = per_batch_noise_counts[ni][batch_idx];
1525 if cnt == 0 {
1526 return None;
1527 }
1528 let mut b = UInt64Builder::with_capacity(cnt);
1529 for _ in 0..cnt {
1530 b.append_value(batch_rng.next_usize(batch_n) as u64);
1531 }
1532 Some((
1533 b.finish(),
1534 batch_rng.next_u64(),
1535 n.noise_type.as_str(),
1536 n.columns.as_slice(),
1537 cnt,
1538 ))
1539 })
1540 .collect();
1541
1542 let all_types: Vec<&str> = plan
1546 .noise_types
1547 .iter()
1548 .map(|n| n.noise_type.as_str())
1549 .collect();
1550 let all_types_ref = &all_types;
1551 let noise_passes = config.noise_passes.max(1);
1552
1553 type DupNoiseResult = Result<(RecordBatch, Vec<usize>, Vec<bool>), String>;
1554 let fk_exclude_ref = &fk_exclude_cols;
1555 let rb_ref = &rb;
1556 let results: Vec<DupNoiseResult> = ndata
1569 .par_iter()
1570 .map(|(indices, seed, ntype, cols, cnt)| {
1571 let dup = pick_rows(rb_ref, indices)?;
1572 let mut rng = Rng::new(*seed);
1573
1574 let (mut noisy, mut unchanged, _) =
1580 apply_noise_with_retry(&dup, ntype, cols, &mut rng, fk_exclude_ref)?;
1581
1582 for _ in 1..noise_passes {
1594 noisy = apply_extra_pass_per_row(
1595 &noisy,
1596 &mut unchanged,
1597 all_types_ref,
1598 &mut rng,
1599 fk_exclude_ref,
1600 )?;
1601 }
1602
1603 let mut mb = Vec::with_capacity(*cnt);
1604 for j in 0..*cnt {
1605 let global_idx = master_base + offset + indices.value(j) as usize;
1610 mb.push(global_idx);
1611 }
1612 Ok((noisy, mb, unchanged))
1613 })
1614 .collect();
1615 let mut dup_is_identical_buf: Vec<bool> = Vec::new();
1616 for res in results {
1617 let (rb, mb, unchanged) = res?;
1618 dup_batches.push(rb);
1619 dup_master_idx_buf.extend(mb);
1620 dup_is_identical_buf.extend(unchanged);
1621 }
1622 _t_dup += t_d0.elapsed().as_secs_f64();
1623
1624 if !dup_batches.is_empty() {
1626 let dup_total = dup_master_idx_buf.len();
1627 let dup_rid_arr =
1628 record_id_array(global_rid_offset..global_rid_offset + dup_total);
1629 let dup_mid_arr = master_id_array_from_indices(
1630 &prefix,
1631 dup_master_idx_buf.iter().copied(),
1632 dup_total,
1633 );
1634
1635 let concated = if dup_batches.len() == 1 {
1636 dup_batches.into_iter().next().unwrap()
1637 } else {
1638 let schema = dup_batches[0].schema();
1639 let n_fields = dup_batches[0].num_columns();
1640 let mut concat_arrays = Vec::with_capacity(n_fields);
1641 for i in 0..n_fields {
1642 let refs: Vec<&dyn Array> =
1643 dup_batches.iter().map(|b| b.column(i).as_ref()).collect();
1644 concat_arrays.push(
1645 arrow::compute::concat(&refs)
1646 .map_err(|e| format!("concat dup col {i}: {e}"))?,
1647 );
1648 }
1649 RecordBatch::try_new(schema, concat_arrays)
1650 .map_err(|e| format!("concat dups: {e}"))?
1651 };
1652
1653 let dup_rb_full = add_metadata_and_align(
1654 &concated,
1655 &config.domain,
1656 &plan.name,
1657 &domain_dict,
1658 &entity_type_dict,
1659 &dup_rid_arr,
1660 &dup_mid_arr,
1661 &full_arc,
1662 col_lookup.as_ref().unwrap(),
1663 &mut null_cache,
1664 &mut const_arr_cache,
1665 )?;
1666
1667 let t_wd = std::time::Instant::now();
1668 writer
1669 .write(&dup_rb_full)
1670 .map_err(|e| format!("write dups: {e}"))?;
1671 _t_write += t_wd.elapsed().as_secs_f64();
1672 _write_calls += 1;
1673
1674 if config.graph_enabled
1676 && let Some(nw) = node_writer.as_mut()
1677 {
1678 nw.write_batch(&dup_rb_full)
1679 .map_err(|e| format!("write dup node: {e}"))?;
1680 }
1681
1682 let is_identical_arr: ArrayRef =
1683 Arc::new(arrow::array::BooleanArray::from(dup_is_identical_buf));
1684 gt_acc.push_dup_batch(
1685 dup_rb_full.column(0),
1686 dup_rb_full.column(2),
1687 dup_rb_full.column(3),
1688 &is_identical_arr,
1689 )?;
1690 global_rid_offset += dup_total;
1691 }
1692 }
1693 }
1694
1695 if let Some(id_col) = &plan.identifier_col
1697 && fk_builder.len() > 0
1698 {
1699 let arr = Arc::new(fk_builder.finish()) as ArrayRef;
1700 if config.graph_enabled {
1701 let rid_arr = Arc::new(fk_rid_builder.finish()) as ArrayRef;
1705 let schema = Arc::new(Schema::new(vec![
1706 Field::new(id_col, DataType::Utf8, true),
1707 Field::new("record_id", DataType::Utf8, true),
1708 ]));
1709 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr, rid_arr]) {
1710 fk_pools.insert(plan.name.clone(), pool_rb);
1711 }
1712 } else {
1713 let schema = Arc::new(Schema::new(vec![Field::new(id_col, DataType::Utf8, true)]));
1714 if let Ok(pool_rb) = RecordBatch::try_new(schema, vec![arr]) {
1715 fk_pools.insert(plan.name.clone(), pool_rb);
1716 }
1717 }
1718 }
1719
1720 if !hn_slices.is_empty() {
1722 let total_count: usize = hn_slices.iter().map(|s| s.num_rows()).sum();
1723 let batch = if hn_slices.len() == 1 {
1724 let only = hn_slices.into_iter().next().unwrap();
1731 let idx = UInt64Array::from_iter_values(0..only.num_rows() as u64);
1732 pick_rows(&only, &idx)?
1733 } else {
1734 let schema = hn_slices[0].schema();
1735 let n_fields = schema.fields().len();
1736 let mut concat_arrays = Vec::with_capacity(n_fields);
1737 for i in 0..n_fields {
1738 let refs: Vec<&dyn Array> =
1739 hn_slices.iter().map(|pb| pb.column(i).as_ref()).collect();
1740 concat_arrays.push(
1741 arrow::compute::concat(&refs)
1742 .map_err(|e| format!("hn pool concat col {i}: {e}"))?,
1743 );
1744 }
1745 RecordBatch::try_new(schema, concat_arrays)
1746 .map_err(|e| format!("hn pool concat batch: {e}"))?
1747 };
1748 let record_ids: ArrayRef = if config.graph_enabled {
1750 let refs: Vec<&dyn Array> = hn_slice_rids.iter().map(|a| a.as_ref()).collect();
1751 arrow::compute::concat(&refs).map_err(|e| format!("hn rid concat: {e}"))?
1752 } else {
1753 Arc::new(StringArray::new_null(0))
1754 };
1755 hn_pools.insert(
1756 plan.name.clone(),
1757 HnPool {
1758 batch,
1759 total_count,
1760 record_ids,
1761 },
1762 );
1763 }
1764 }
1765 let t1e_elapsed = t1e.elapsed().as_secs_f64();
1766 log::debug!("[sink_profile] write={_t_write:.3}s calls={_write_calls}");
1767 log_rss("after entity batches (fk_pools+hn_pools built)");
1768
1769 let t2 = std::time::Instant::now();
1771 let mut hn_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1773 let mut hn_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1774 let mut hn_master_id_counter: u64 = offsets.hn_master_counter;
1782 for hn_cfg in config.hard_neg_types.iter() {
1783 if hn_cfg.count == 0 {
1784 continue;
1785 }
1786 let Some(pool_data) = hn_pools.get(&hn_cfg.entity_type) else {
1787 continue;
1788 };
1789 let hn_total = pool_data.total_count;
1790
1791 let pool_rb = &pool_data.batch;
1792 let n_hn_max = hn_cfg.count.min(hn_total / 2);
1793 if n_hn_max == 0 {
1794 continue;
1795 }
1796
1797 let (hn_rb, hn_src) = crate::hn_common::generate_hard_negatives(
1798 pool_rb,
1799 &hn_cfg.config_json,
1800 n_hn_max,
1801 config.seed.wrapping_add(200),
1802 config.graph_enabled,
1803 )?;
1804
1805 let n_hn = hn_rb.num_rows();
1806 if n_hn == 0 {
1807 continue;
1808 }
1809
1810 let hn_rid_arr = record_id_array(global_rid_offset..global_rid_offset + n_hn);
1811 let hn_mid_arr = hn_master_id_array(hn_master_id_counter, n_hn);
1812 hn_master_id_counter += n_hn as u64;
1813
1814 let hn_schema = hn_rb.schema();
1816 let hn_col_lookup: Vec<Option<usize>> = full_arc
1817 .fields()
1818 .iter()
1819 .skip(4)
1820 .map(|f| hn_schema.column_with_name(f.name()).map(|(idx, _)| idx))
1821 .collect();
1822
1823 let hn_rb_full = add_metadata_and_align(
1824 &hn_rb,
1825 &config.domain,
1826 &hn_cfg.entity_type,
1827 &domain_dict,
1828 &entity_type_dict,
1829 &hn_rid_arr,
1830 &hn_mid_arr,
1831 &full_arc,
1832 &hn_col_lookup,
1833 &mut hn_null_cache,
1834 &mut hn_const_arr_cache,
1835 )?;
1836
1837 let t_wh = std::time::Instant::now();
1838 writer
1839 .write(&hn_rb_full)
1840 .map_err(|e| format!("write hn: {e}"))?;
1841 _t_write += t_wh.elapsed().as_secs_f64();
1842 _write_calls += 1;
1843
1844 if config.graph_enabled
1847 && let (Some(nw), Some(ew)) = (node_writer.as_mut(), edge_writer.as_mut())
1848 {
1849 nw.write_batch(&hn_rb_full)
1850 .map_err(|e| format!("write hn node: {e}"))?;
1851 if let Some((idx_a, pattern)) = hn_src {
1852 let hn_rid_str_arr = hn_rid_arr.as_string::<i32>();
1853 let pool_rid_arr = pool_data.record_ids.as_string::<i32>();
1854 for (i, &pool_idx) in idx_a.iter().enumerate().take(n_hn) {
1855 let tgt = pool_rid_arr.value(pool_idx);
1856 ew.push(hn_rid_str_arr.value(i), tgt, "hard_neg", &pattern, 1.0)
1857 .map_err(|e| format!("push hn edge: {e}"))?;
1858 }
1859 }
1860 }
1861
1862 gt_acc.push_other_batch(
1864 hn_rb_full.column(0),
1865 hn_rb_full.column(2),
1866 hn_rb_full.column(3),
1867 )?;
1868
1869 global_rid_offset += n_hn;
1870 }
1871 let t2_elapsed = t2.elapsed().as_secs_f64();
1872 drop(hn_pools);
1879 log_rss("after hard negatives");
1880
1881 {
1883 let mut canary_null_cache: FxHashMap<(DataType, usize), ArrayRef> = FxHashMap::default();
1884 let mut canary_const_arr_cache: FxHashMap<(String, usize), ArrayRef> = FxHashMap::default();
1885 crate::canary::generate_all(
1886 ctx,
1887 config,
1888 &full_arc,
1889 &domain_dict,
1890 &entity_type_dict,
1891 &mut canary_null_cache,
1892 &mut canary_const_arr_cache,
1893 &mut global_rid_offset,
1894 &fk_pools,
1895 &mut writer,
1896 &mut node_writer,
1897 &mut gt_acc,
1898 )?;
1899 }
1900 log_rss("after canary records");
1901
1902 let t3 = std::time::Instant::now();
1904 writer
1905 .finish()
1906 .map_err(|e| format!("finish dataset writer: {e}"))?;
1907 let t3a_elapsed = t3.elapsed().as_secs_f64();
1908 log_rss("after dataset writer finish");
1909
1910 let t3b = std::time::Instant::now();
1911 let t_gt0 = std::time::Instant::now();
1912 let crate::gt::GtResult {
1913 n_exact_dup,
1914 n_fuzzy_dup,
1915 n_hard_neg,
1916 n_unique,
1917 n_masters,
1918 cluster_map,
1919 } = gt_acc.finish(
1920 config.difficulty.as_str(),
1921 &config.output_format,
1922 >_path,
1923 &metadata,
1924 config.graph_enabled,
1925 )?;
1926 let _t_gt_compute = 0.0f64;
1927 let _t_gt_write = t_gt0.elapsed().as_secs_f64();
1928 let t3b_elapsed = t3b.elapsed().as_secs_f64();
1929 let _t3_elapsed = t3.elapsed().as_secs_f64();
1930 log_rss("after GT finish (cluster_map built)");
1931
1932 if config.graph_enabled
1934 && let Some(ew) = edge_writer.as_mut()
1935 {
1936 crate::graph_gen::push_dup_clusters(ew, &cluster_map, GRAPH_MAX_CLUSTER_EDGES)
1937 .map_err(|e| format!("push dup clusters: {e}"))?;
1938 }
1939
1940 let duration = t_start.elapsed().as_secs_f64();
1941
1942 log::debug!(
1943 "[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",
1944 t_alloc,
1945 t1e_elapsed,
1946 t2_elapsed,
1947 t3a_elapsed,
1948 _t_meta,
1949 _t_dup,
1950 _t_write,
1951 _t_gt_compute,
1952 _t_gt_write,
1953 t3b_elapsed,
1954 duration,
1955 );
1956
1957 let stats = PipelineStats {
1958 total_records: global_rid_offset,
1959 exact_dups: n_exact_dup,
1960 fuzzy_dups: n_fuzzy_dup,
1961 hard_negs: n_hard_neg,
1962 uniques: n_unique,
1963 masters: n_masters,
1964 };
1965
1966 let mut graph_nodes_final: Option<String> = None;
1968 let mut graph_edges_final: Option<String> = None;
1969 if let (Some(nw), Some(ew)) = (node_writer, edge_writer) {
1970 nw.finish()
1971 .map_err(|e| format!("finish node writer: {e}"))?;
1972 ew.finish()
1973 .map_err(|e| format!("finish edge writer: {e}"))?;
1974 let mut np = graph_nodes_path.clone().unwrap();
1975 let mut ep = graph_edges_path.clone().unwrap();
1976 if graph_fmt == crate::graph_gen::GraphFormat::Parquet {
1977 let np_pq = format!("{}/{}_nodes.parquet", output_dir, config.run_id);
1978 let ep_pq = format!("{}/{}_edges.parquet", output_dir, config.run_id);
1979 convert_ipc_to_parquet(&np, &np_pq).map_err(|e| format!("graph nodes parquet: {e}"))?;
1980 convert_ipc_to_parquet(&ep, &ep_pq).map_err(|e| format!("graph edges parquet: {e}"))?;
1981 std::fs::remove_file(&np).ok();
1982 std::fs::remove_file(&ep).ok();
1983 np = np_pq;
1984 ep = ep_pq;
1985 }
1986 graph_nodes_final = Some(np);
1987 graph_edges_final = Some(ep);
1988 }
1989 log_rss("end (after graph finalize)");
1990
1991 let mut master_local = offsets.master_local.clone();
1995 if master_local.len() < config.entity_plans.len() {
1996 master_local.resize(config.entity_plans.len(), 0);
1997 }
1998 for (idx, plan) in config.entity_plans.iter().enumerate() {
1999 if !plan.pool_only {
2000 master_local[idx] = offsets.master_base(idx) + plan.n_base;
2001 }
2002 }
2003 let next_offsets = ChunkOffsets {
2004 rid_offset: global_rid_offset,
2005 master_local,
2006 hn_master_counter: hn_master_id_counter,
2007 };
2008
2009 Ok((
2010 PipelineOutput {
2011 output_files: vec![dataset_path],
2012 gt_file: gt_path,
2013 stats,
2014 nodes: graph_nodes_final,
2015 edges: graph_edges_final,
2016 },
2017 next_offsets,
2018 ))
2019}
2020
2021fn convert_ipc_to_parquet(ipc_path: &str, parquet_path: &str) -> Result<(), String> {
2025 use arrow::ipc::reader::FileReader;
2026
2027 let ipc_file =
2028 std::fs::File::open(ipc_path).map_err(|e| format!("open IPC {ipc_path}: {e}"))?;
2029 let ipc_reader =
2030 FileReader::try_new(ipc_file, None).map_err(|e| format!("IPC reader {ipc_path}: {e}"))?;
2031 let schema = ipc_reader.schema();
2032
2033 let parquet_file =
2034 std::fs::File::create(parquet_path).map_err(|e| format!("create {parquet_path}: {e}"))?;
2035 let zstd_level = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
2036 let meta_kv: Vec<parquet::file::metadata::KeyValue> = schema
2037 .metadata()
2038 .iter()
2039 .map(|(k, v)| parquet::file::metadata::KeyValue {
2040 key: k.clone(),
2041 value: Some(v.clone()),
2042 })
2043 .collect();
2044 let props = parquet::file::properties::WriterProperties::builder()
2045 .set_compression(parquet::basic::Compression::ZSTD(zstd_level))
2046 .set_max_row_group_row_count(Some(1_000_000))
2047 .set_key_value_metadata(Some(meta_kv))
2048 .build();
2049 let mut parquet_writer =
2050 parquet::arrow::ArrowWriter::try_new(parquet_file, schema, Some(props))
2051 .map_err(|e| format!("ArrowWriter {parquet_path}: {e}"))?;
2052
2053 for batch_result in ipc_reader {
2054 let batch = batch_result.map_err(|e| format!("IPC batch {ipc_path}: {e}"))?;
2055 parquet_writer
2056 .write(&batch)
2057 .map_err(|e| format!("write parquet {parquet_path}: {e}"))?;
2058 }
2059 parquet_writer
2060 .close()
2061 .map_err(|e| format!("close parquet {parquet_path}: {e}"))?;
2062 Ok(())
2063}
2064
2065fn concat_dataset_files(
2074 output_format: &str,
2075 in_paths: &[String],
2076 out_path: &str,
2077) -> Result<(), String> {
2078 if in_paths.is_empty() {
2079 return Err("concat_dataset_files: no chunk files to concatenate".into());
2080 }
2081 if output_format == "parquet" {
2082 concat_parquet_files(in_paths, out_path)
2083 } else {
2084 concat_ipc_files(in_paths, out_path)
2085 }
2086}
2087
2088fn concat_ipc_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
2089 use arrow::ipc::reader::FileReader;
2090
2091 let first_file =
2092 std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
2093 let schema = FileReader::try_new(first_file, None)
2094 .map_err(|e| format!("ipc reader {}: {e}", in_paths[0]))?
2095 .schema();
2096
2097 let out_file =
2098 std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
2099 let mut writer = arrow::ipc::writer::FileWriter::try_new(out_file, &schema)
2100 .map_err(|e| format!("ipc writer {out_path}: {e}"))?;
2101
2102 for path in in_paths {
2103 let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
2104 let reader =
2105 FileReader::try_new(file, None).map_err(|e| format!("ipc reader {path}: {e}"))?;
2106 for batch_result in reader {
2107 let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
2108 writer
2109 .write(&batch)
2110 .map_err(|e| format!("write {out_path}: {e}"))?;
2111 }
2112 }
2113 writer
2114 .finish()
2115 .map_err(|e| format!("finish {out_path}: {e}"))
2116}
2117
2118fn concat_parquet_files(in_paths: &[String], out_path: &str) -> Result<(), String> {
2119 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
2120
2121 let first_file =
2122 std::fs::File::open(&in_paths[0]).map_err(|e| format!("open {}: {e}", in_paths[0]))?;
2123 let first_builder = ParquetRecordBatchReaderBuilder::try_new(first_file)
2124 .map_err(|e| format!("parquet reader {}: {e}", in_paths[0]))?;
2125 let schema = first_builder.schema().clone();
2126 let meta_kv: Vec<parquet::file::metadata::KeyValue> = first_builder
2127 .metadata()
2128 .file_metadata()
2129 .key_value_metadata()
2130 .cloned()
2131 .unwrap_or_default();
2132
2133 let out_file =
2134 std::fs::File::create(out_path).map_err(|e| format!("create {out_path}: {e}"))?;
2135 let zstd = parquet::basic::ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
2136 let props = parquet::file::properties::WriterProperties::builder()
2137 .set_compression(parquet::basic::Compression::ZSTD(zstd))
2138 .set_max_row_group_row_count(Some(1_000_000))
2139 .set_key_value_metadata(Some(meta_kv))
2140 .build();
2141 let mut writer = parquet::arrow::ArrowWriter::try_new(out_file, schema, Some(props))
2142 .map_err(|e| format!("arrow writer {out_path}: {e}"))?;
2143
2144 for path in in_paths {
2145 let file = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
2146 let reader = ParquetRecordBatchReaderBuilder::try_new(file)
2147 .map_err(|e| format!("parquet reader {path}: {e}"))?
2148 .build()
2149 .map_err(|e| format!("build parquet reader {path}: {e}"))?;
2150 for batch_result in reader {
2151 let batch = batch_result.map_err(|e| format!("read {path}: {e}"))?;
2152 writer
2153 .write(&batch)
2154 .map_err(|e| format!("write {out_path}: {e}"))?;
2155 }
2156 }
2157 writer
2158 .close()
2159 .map(|_| ())
2160 .map_err(|e| format!("close {out_path}: {e}"))
2161}
2162
2163#[allow(clippy::too_many_arguments)]
2177pub fn run_chunked(
2178 ctx: &Context,
2179 domain: &str,
2180 total_size: usize,
2181 chunk_size: usize,
2182 seed: u64,
2183 difficulty: &str,
2184 hard_neg_ratio: f64,
2185 singleton_master_fraction: f64,
2186 schema: &crate::schema::DomainSchema,
2187 final_run_id: &str,
2188 output_format: &str,
2189 graph_enabled: bool,
2190 graph_format: &str,
2191 only_entity: Option<&str>,
2192 output_dir: &str,
2193 mut progress: Option<&mut dyn FnMut(usize, usize)>,
2194) -> Result<PipelineOutput, String> {
2195 if chunk_size == 0 || chunk_size >= total_size {
2196 let config = crate::schema::build_pipeline_config(
2197 domain,
2198 total_size,
2199 seed,
2200 difficulty,
2201 hard_neg_ratio,
2202 singleton_master_fraction,
2203 schema,
2204 final_run_id,
2205 output_format,
2206 graph_enabled,
2207 graph_format,
2208 only_entity,
2209 )?;
2210 return run_pipeline_with_progress(ctx, &config, output_dir, progress);
2211 }
2212
2213 let n_chunks = total_size.div_ceil(chunk_size);
2214 let tmp_dir = format!("{output_dir}/.dupehell_chunks_{final_run_id}");
2215 std::fs::create_dir_all(&tmp_dir).map_err(|e| format!("create {tmp_dir:?}: {e}"))?;
2216
2217 let mut offsets = ChunkOffsets::default();
2218 let mut dataset_chunks: Vec<String> = Vec::new();
2219 let mut gt_chunks: Vec<String> = Vec::new();
2220 let mut node_chunks: Vec<String> = Vec::new();
2221 let mut edge_chunks: Vec<String> = Vec::new();
2222 let mut agg = PipelineStats {
2223 total_records: 0,
2224 exact_dups: 0,
2225 fuzzy_dups: 0,
2226 hard_negs: 0,
2227 uniques: 0,
2228 masters: 0,
2229 };
2230
2231 for chunk_idx in 0..n_chunks {
2232 let this_size = chunk_size.min(total_size - chunk_idx * chunk_size);
2233 let chunk_seed = seed.wrapping_add(chunk_idx as u64);
2234 let chunk_run_id = format!("{final_run_id}__chunk{chunk_idx:04}");
2235 let chunk_config = crate::schema::build_pipeline_config(
2236 domain,
2237 this_size,
2238 chunk_seed,
2239 difficulty,
2240 hard_neg_ratio,
2241 singleton_master_fraction,
2242 schema,
2243 &chunk_run_id,
2244 output_format,
2245 graph_enabled,
2246 graph_format,
2247 only_entity,
2248 )?;
2249
2250 let already_written = agg.total_records;
2251 let mut chunk_cb = |done: usize, _total: usize| {
2252 if let Some(cb) = &mut progress {
2253 cb(already_written + done, total_size);
2254 }
2255 };
2256 let (out, next_offsets) = run_pipeline_chunked(
2257 ctx,
2258 &chunk_config,
2259 &tmp_dir,
2260 Some(&mut chunk_cb),
2261 Some(&offsets),
2262 )?;
2263 offsets = next_offsets;
2264
2265 agg.total_records = out.stats.total_records;
2270 agg.exact_dups += out.stats.exact_dups;
2271 agg.fuzzy_dups += out.stats.fuzzy_dups;
2272 agg.hard_negs += out.stats.hard_negs;
2273 agg.uniques += out.stats.uniques;
2274 agg.masters += out.stats.masters;
2275
2276 dataset_chunks.push(out.output_files[0].clone());
2277 gt_chunks.push(out.gt_file.clone());
2278 if let Some(n) = out.nodes {
2279 node_chunks.push(n);
2280 }
2281 if let Some(e) = out.edges {
2282 edge_chunks.push(e);
2283 }
2284 }
2285
2286 let dataset_ext = if output_format == "parquet" {
2287 "parquet"
2288 } else {
2289 "ipc"
2290 };
2291 let final_dataset = format!("{output_dir}/{final_run_id}.{dataset_ext}");
2292 concat_dataset_files(output_format, &dataset_chunks, &final_dataset)?;
2293
2294 let final_gt = format!("{output_dir}/{final_run_id}_ground_truth.{dataset_ext}");
2295 concat_dataset_files(output_format, >_chunks, &final_gt)?;
2296
2297 let graph_ext = if graph_format == "parquet" {
2298 "parquet"
2299 } else {
2300 "ipc"
2301 };
2302 let final_nodes = if !node_chunks.is_empty() {
2303 let p = format!("{output_dir}/{final_run_id}_nodes.{graph_ext}");
2304 concat_dataset_files(graph_format, &node_chunks, &p)?;
2305 Some(p)
2306 } else {
2307 None
2308 };
2309 let final_edges = if !edge_chunks.is_empty() {
2310 let p = format!("{output_dir}/{final_run_id}_edges.{graph_ext}");
2311 concat_dataset_files(graph_format, &edge_chunks, &p)?;
2312 Some(p)
2313 } else {
2314 None
2315 };
2316
2317 std::fs::remove_dir_all(&tmp_dir).ok();
2318
2319 Ok(PipelineOutput {
2320 output_files: vec![final_dataset],
2321 gt_file: final_gt,
2322 stats: agg,
2323 nodes: final_nodes,
2324 edges: final_edges,
2325 })
2326}
2327
2328fn build_metadata_map(config: &PipelineConfig) -> HashMap<String, String> {
2331 let ts = std::time::SystemTime::now()
2332 .duration_since(std::time::UNIX_EPOCH)
2333 .map(|d| d.as_secs().to_string())
2334 .unwrap_or_else(|_| "unknown".into());
2335 HashMap::from([
2336 (
2337 "dupehell.generator".into(),
2338 format!("DupeHell v{}", env!("CARGO_PKG_VERSION")),
2339 ),
2340 (
2341 "dupehell.provenance".into(),
2342 "dupehell-synthetic-data".into(),
2343 ),
2344 ("dupehell.license".into(), "MIT".into()),
2345 (
2346 "dupehell.purpose".into(),
2347 "Educational Use Only — Record Linkage Benchmarking".into(),
2348 ),
2349 (
2350 "dupehell.url".into(),
2351 "https://github.com/vntoinekaio/DupeHell".into(),
2352 ),
2353 ("dupehell.domain".into(), config.domain.clone()),
2354 ("dupehell.size".into(), config.size.to_string()),
2355 ("dupehell.seed".into(), config.seed.to_string()),
2356 ("dupehell.difficulty".into(), config.difficulty.clone()),
2357 ("dupehell.run_id".into(), config.run_id.clone()),
2358 ("dupehell.timestamp".into(), ts),
2359 (
2360 "dupehell.structural_key_columns".into(),
2361 structural_key_columns_note(config),
2362 ),
2363 ])
2364}
2365
2366fn structural_key_columns_note(config: &PipelineConfig) -> String {
2371 let cols: Vec<&str> = config
2372 .entity_plans
2373 .iter()
2374 .filter_map(|p| p.identifier_col.as_deref())
2375 .collect();
2376 format!(
2377 "{} — structural keys for joining child tables, not ER match attributes; use record_id as the row identifier instead",
2378 cols.join(", ")
2379 )
2380}
2381
2382pub(crate) fn low_cardinality_dict_type() -> DataType {
2391 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
2392}
2393
2394pub(crate) struct DictValues {
2411 values: ArrayRef,
2412 index: HashMap<String, i32>,
2413}
2414
2415impl DictValues {
2416 pub(crate) fn new<I, S>(values: I) -> Self
2417 where
2418 I: IntoIterator<Item = S>,
2419 S: Into<String>,
2420 {
2421 let owned: Vec<String> = values.into_iter().map(Into::into).collect();
2422 let index: HashMap<String, i32> = owned
2423 .iter()
2424 .enumerate()
2425 .map(|(i, v)| (v.clone(), i as i32))
2426 .collect();
2427 Self {
2428 values: Arc::new(StringArray::from(owned)),
2429 index,
2430 }
2431 }
2432
2433 fn idx(&self, value: &str) -> i32 {
2434 *self.index.get(value).unwrap_or_else(|| {
2435 panic!("DictValues: value {value:?} was not in the pre-built dictionary")
2436 })
2437 }
2438
2439 fn build(&self, keys: arrow::array::Int32Array) -> ArrayRef {
2440 Arc::new(
2441 arrow::array::DictionaryArray::<arrow::datatypes::Int32Type>::try_new(
2442 keys,
2443 self.values.clone(),
2444 )
2445 .expect("DictValues::build: key array out of range for this dictionary"),
2446 )
2447 }
2448
2449 pub(crate) fn const_array(&self, value: &str, n: usize) -> ArrayRef {
2451 self.build(arrow::array::Int32Array::from(vec![self.idx(value); n]))
2452 }
2453
2454 pub(crate) fn key(&self, value: &str) -> i32 {
2459 self.idx(value)
2460 }
2461
2462 pub(crate) fn finish_keys(&self, keys: arrow::array::Int32Array) -> ArrayRef {
2463 self.build(keys)
2464 }
2465}
2466
2467fn subtype_dict(config: &PipelineConfig) -> DictValues {
2477 let mut values: Vec<String> = vec!["complete".to_string(), "spanning_tree".to_string()];
2478 for plan in &config.entity_plans {
2479 for remap in &plan.fk_remaps {
2480 values.push(remap.source_col.clone());
2481 }
2482 }
2483 for hn in &config.hard_neg_types {
2484 if let Ok(v) = serde_json::from_str::<serde_json::Value>(&hn.config_json)
2485 && let Some(pattern) = v.get("pattern").and_then(|p| p.as_str())
2486 {
2487 values.push(pattern.to_string());
2488 }
2489 }
2490 values.sort_unstable();
2491 values.dedup();
2492 DictValues::new(values)
2493}
2494
2495fn build_full_schema(config: &PipelineConfig, metadata: &HashMap<String, String>) -> Schema {
2496 let mut field_map: Vec<(String, DataType, bool)> = vec![
2497 ("record_id".to_string(), DataType::Utf8, false),
2498 ("domain".to_string(), low_cardinality_dict_type(), false),
2499 (
2500 "entity_type".to_string(),
2501 low_cardinality_dict_type(),
2502 false,
2503 ),
2504 ("master_id".to_string(), DataType::Utf8, false),
2505 ];
2506 for plan in &config.entity_plans {
2507 if plan.pool_only {
2512 continue;
2513 }
2514 let cols: Vec<serde_json::Value> =
2515 serde_json::from_str(&plan.columns_json).unwrap_or_default();
2516 for col in &cols {
2517 let name = col["name"].as_str().unwrap_or("").to_string();
2518 if name.is_empty() || field_map.iter().any(|(n, _, _)| n == &name) {
2519 continue;
2520 }
2521 let col_type = col_type_from_request(col);
2522 field_map.push((name, col_type, true));
2523 }
2524 }
2525 let fields: Vec<Field> = field_map
2526 .into_iter()
2527 .map(|(n, dt, nullable)| Field::new(&n, dt, nullable))
2528 .collect();
2529 Schema::new(fields).with_metadata(metadata.clone())
2530}
2531
2532fn col_type_from_request(v: &serde_json::Value) -> DataType {
2534 match v["type"].as_str() {
2535 Some("int") => DataType::Int64,
2536 Some("float") => DataType::Float64,
2537 Some("boolean") => DataType::Boolean,
2538 Some("date") | Some("datetime") => DataType::Utf8,
2539 _ => DataType::Utf8,
2540 }
2541}
2542
2543#[allow(clippy::too_many_arguments)]
2546pub(crate) fn add_metadata_and_align(
2547 rb: &RecordBatch,
2548 domain: &str,
2549 entity_type: &str,
2550 domain_dict: &DictValues,
2551 entity_type_dict: &DictValues,
2552 record_ids: &ArrayRef,
2553 master_ids: &ArrayRef,
2554 full_arc: &Arc<Schema>,
2555 col_lookup: &[Option<usize>],
2556 null_cache: &mut FxHashMap<(DataType, usize), ArrayRef>,
2557 const_arr_cache: &mut FxHashMap<(String, usize), ArrayRef>,
2558) -> Result<RecordBatch, String> {
2559 let n = rb.num_rows();
2560 let domain_arr = const_arr_cache
2570 .entry((domain.to_string(), n))
2571 .or_insert_with(|| domain_dict.const_array(domain, n))
2572 .clone();
2573 let et_arr = const_arr_cache
2574 .entry((entity_type.to_string(), n))
2575 .or_insert_with(|| entity_type_dict.const_array(entity_type, n))
2576 .clone();
2577 let rid_arr = record_ids.clone();
2578 let mid_arr = master_ids.clone();
2579
2580 let mut all_arrays: Vec<ArrayRef> = Vec::with_capacity(4 + col_lookup.len());
2581 all_arrays.push(rid_arr);
2582 all_arrays.push(domain_arr);
2583 all_arrays.push(et_arr);
2584 all_arrays.push(mid_arr);
2585
2586 for (i, maybe_idx) in col_lookup.iter().enumerate() {
2589 match maybe_idx {
2590 Some(idx) => all_arrays.push(rb.column(*idx).clone()),
2591 None => {
2592 let dt = full_arc.field(i + 4).data_type();
2593 let arr = null_cache
2594 .entry((dt.clone(), n))
2595 .or_insert_with(|| arrow::array::new_null_array(dt, n));
2596 all_arrays.push(Arc::clone(arr));
2597 }
2598 }
2599 }
2600
2601 RecordBatch::try_new(Arc::clone(full_arc), all_arrays)
2603 .map_err(|e| format!("add_metadata_and_align RecordBatch: {e}"))
2604}
2605
2606fn pick_rows(rb: &RecordBatch, indices: &UInt64Array) -> Result<RecordBatch, String> {
2607 use arrow::compute::take;
2608 let new_columns: Vec<ArrayRef> = (0..rb.num_columns())
2609 .map(|i| take(rb.column(i), indices, None))
2610 .collect::<Result<Vec<_>, _>>()
2611 .map_err(|e| format!("take error: {e}"))?;
2612
2613 RecordBatch::try_new(rb.schema(), new_columns).map_err(|e| format!("RecordBatch: {e}"))
2614}
2615
2616#[cfg(test)]
2617mod tests {
2618 use super::*;
2619 use crate::schema::{build_pipeline_config, load_schema};
2620
2621 #[test]
2628 fn test_no_person_company_name_cross_contamination() {
2629 for col in [
2630 "given_name",
2631 "family_name",
2632 "middle_name",
2633 "first_name",
2634 "last_name",
2635 ] {
2636 assert!(
2637 noise_type_targets_column("names", col),
2638 "{col} should still match 'names'"
2639 );
2640 assert!(
2641 !noise_type_targets_column("companies", col),
2642 "{col} should NOT match 'companies' (person name, not a company name)"
2643 );
2644 }
2645 for col in ["company_name", "legal_name", "trading_name"] {
2646 assert!(
2647 noise_type_targets_column("companies", col),
2648 "{col} should still match 'companies'"
2649 );
2650 assert!(
2651 !noise_type_targets_column("names", col),
2652 "{col} should NOT match 'names' (company name, not a person name)"
2653 );
2654 }
2655 assert!(
2656 !noise_type_targets_column("addresses", "email_address"),
2657 "email_address should NOT match 'addresses' (it's an email, not a postal address)"
2658 );
2659 assert!(
2660 noise_type_targets_column("identifiers", "email_address"),
2661 "email_address should still match 'identifiers' (corrupt_email)"
2662 );
2663 assert!(
2664 noise_type_targets_column("addresses", "residential_address_line1"),
2665 "a real postal address column should still match 'addresses'"
2666 );
2667 }
2668
2669 fn manifest_dir() -> std::path::PathBuf {
2670 std::path::Path::new(env!("CARGO_MANIFEST_DIR")).to_path_buf()
2671 }
2672
2673 fn scratch_dir(name: &str) -> std::path::PathBuf {
2676 let dir = manifest_dir()
2677 .join("target")
2678 .join("pipeline_test_tmp")
2679 .join(format!("{name}_{}", std::process::id()));
2680 let _ = std::fs::remove_dir_all(&dir);
2681 std::fs::create_dir_all(&dir).expect("create scratch dir");
2682 dir
2683 }
2684
2685 fn run_kyc(
2686 output_dir: &std::path::Path,
2687 output_format: &str,
2688 graph_enabled: bool,
2689 ) -> PipelineOutput {
2690 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2691 let mut ctx = Context::new(
2692 "kyc",
2693 "en",
2694 &manifest_dir().join("assets/pools").to_string_lossy(),
2695 )
2696 .expect("load context");
2697 let config = build_pipeline_config(
2698 "kyc",
2699 200,
2700 42,
2701 "medium",
2702 0.1,
2703 0.3,
2704 &schema,
2705 "pipeline_test_run",
2706 output_format,
2707 graph_enabled,
2708 "parquet",
2709 None,
2710 )
2711 .expect("build config");
2712 ctx.enable_watermark(&config.domain, config.size, config.seed);
2713 run_pipeline(&ctx, &config, &output_dir.to_string_lossy()).expect("run_pipeline")
2714 }
2715
2716 #[test]
2717 fn test_run_pipeline_produces_dataset_and_gt() {
2718 let dir = scratch_dir("basic");
2719 let output = run_kyc(&dir, "parquet", false);
2720
2721 assert_eq!(output.output_files.len(), 1);
2722 for f in &output.output_files {
2723 assert!(std::path::Path::new(f).exists(), "missing dataset file {f}");
2724 }
2725 assert!(std::path::Path::new(&output.gt_file).exists());
2726 assert!(output.nodes.is_none());
2727 assert!(output.edges.is_none());
2728
2729 let s = &output.stats;
2734 assert!(s.total_records >= s.exact_dups + s.fuzzy_dups + s.hard_negs + s.uniques);
2735 assert!(s.total_records > 0);
2736 assert!(s.masters > 0);
2737
2738 let _ = std::fs::remove_dir_all(&dir);
2739 }
2740
2741 #[test]
2742 fn test_run_pipeline_ipc_format() {
2743 let dir = scratch_dir("ipc");
2744 let output = run_kyc(&dir, "ipc", false);
2745 assert!(output.output_files[0].ends_with(".ipc"));
2746 assert!(std::path::Path::new(&output.output_files[0]).exists());
2747 let _ = std::fs::remove_dir_all(&dir);
2748 }
2749
2750 #[test]
2751 fn test_run_pipeline_graph_output() {
2752 let dir = scratch_dir("graph");
2753 let output = run_kyc(&dir, "parquet", true);
2754 let nodes = output.nodes.expect("nodes path");
2755 let edges = output.edges.expect("edges path");
2756 assert!(std::path::Path::new(&nodes).exists());
2757 assert!(std::path::Path::new(&edges).exists());
2758 let _ = std::fs::remove_dir_all(&dir);
2759 }
2760
2761 #[test]
2762 fn test_fk_edge_weight_from_schema_and_default() {
2763 let dir = scratch_dir("fk_weight");
2770 let schema =
2771 load_schema("aviation", &manifest_dir().join("schemas")).expect("load aviation.json");
2772 let mut ctx = Context::new(
2773 "aviation",
2774 "en",
2775 &manifest_dir().join("assets/pools").to_string_lossy(),
2776 )
2777 .expect("load context");
2778 let config = build_pipeline_config(
2779 "aviation",
2780 300,
2781 42,
2782 "medium",
2783 0.1,
2784 0.3,
2785 &schema,
2786 "pipeline_test_fk_weight",
2787 "ipc",
2788 true,
2789 "ipc",
2790 None,
2791 )
2792 .expect("build config");
2793 ctx.enable_watermark(&config.domain, config.size, config.seed);
2794 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2795
2796 let edges_path = output.edges.expect("edges path");
2797 let file = std::fs::File::open(&edges_path).expect("open edges file");
2798 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2799
2800 let mut subtypes_seen: std::collections::HashMap<String, f64> =
2801 std::collections::HashMap::new();
2802 for batch in reader {
2803 let batch = batch.expect("read edges batch");
2804 let edge_type_col = arrow::compute::cast(
2808 batch.column_by_name("edge_type").expect("edge_type col"),
2809 &DataType::Utf8,
2810 )
2811 .expect("cast edge_type to Utf8");
2812 let edge_type = edge_type_col.as_string::<i32>();
2813 let subtype_col = arrow::compute::cast(
2814 batch.column_by_name("subtype").expect("subtype col"),
2815 &DataType::Utf8,
2816 )
2817 .expect("cast subtype to Utf8");
2818 let subtype = subtype_col.as_string::<i32>();
2819 let weight = batch
2820 .column_by_name("weight")
2821 .expect("weight col")
2822 .as_any()
2823 .downcast_ref::<arrow::array::Float64Array>()
2824 .expect("weight as f64");
2825 for i in 0..batch.num_rows() {
2826 if edge_type.value(i) == "fk" {
2827 subtypes_seen.insert(subtype.value(i).to_string(), weight.value(i));
2828 }
2829 }
2830 }
2831
2832 assert_eq!(
2833 subtypes_seen.get("aircraft_id"),
2834 Some(&0.5),
2835 "flight->aircraft fk_weight from schema not propagated to edge weight"
2836 );
2837 assert_eq!(
2838 subtypes_seen.get("airline_id"),
2839 Some(&1.0),
2840 "flight->airline fk_weight (schema default) not propagated to edge weight"
2841 );
2842
2843 let _ = std::fs::remove_dir_all(&dir);
2844 }
2845
2846 #[test]
2853 fn test_run_pipeline_dup_cluster_edges_well_formed() {
2854 let dir = scratch_dir("dup_cluster_edges");
2855 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
2860 let mut ctx = Context::new(
2861 "kyc",
2862 "en",
2863 &manifest_dir().join("assets/pools").to_string_lossy(),
2864 )
2865 .expect("load context");
2866 let config = build_pipeline_config(
2867 "kyc",
2868 200,
2869 42,
2870 "medium",
2871 0.1,
2872 0.3,
2873 &schema,
2874 "pipeline_test_dup_cluster_edges",
2875 "ipc",
2876 true,
2877 "ipc",
2878 None,
2879 )
2880 .expect("build config");
2881 ctx.enable_watermark(&config.domain, config.size, config.seed);
2882 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
2883
2884 let edges_path = output.edges.expect("edges path");
2885 let file = std::fs::File::open(&edges_path).expect("open edges file");
2886 let reader = arrow::ipc::reader::FileReader::try_new(file, None).expect("edges reader");
2887
2888 let mut n_exact = 0usize;
2889 let mut n_fuzzy = 0usize;
2890 for batch in reader {
2891 let batch = batch.expect("read edges batch");
2892 let src = batch
2893 .column_by_name("source_node_id")
2894 .expect("source_node_id col")
2895 .as_any()
2896 .downcast_ref::<arrow::array::StringArray>()
2897 .expect("source_node_id as string");
2898 let tgt = batch
2899 .column_by_name("target_node_id")
2900 .expect("target_node_id col")
2901 .as_any()
2902 .downcast_ref::<arrow::array::StringArray>()
2903 .expect("target_node_id as string");
2904 let edge_type_col = arrow::compute::cast(
2907 batch.column_by_name("edge_type").expect("edge_type col"),
2908 &DataType::Utf8,
2909 )
2910 .expect("cast edge_type to Utf8");
2911 let edge_type = edge_type_col.as_string::<i32>();
2912 let weight = batch
2913 .column_by_name("weight")
2914 .expect("weight col")
2915 .as_any()
2916 .downcast_ref::<arrow::array::Float64Array>()
2917 .expect("weight as f64");
2918
2919 for i in 0..batch.num_rows() {
2920 match edge_type.value(i) {
2921 "exact_dup" => n_exact += 1,
2922 "fuzzy_dup" => n_fuzzy += 1,
2923 _ => continue,
2924 }
2925 assert!(
2930 parse_record_idx(src.value(i)).is_some(),
2931 "malformed source_node_id {:?} on a dup-cluster edge",
2932 src.value(i)
2933 );
2934 assert!(
2935 parse_record_idx(tgt.value(i)).is_some(),
2936 "malformed target_node_id {:?} on a dup-cluster edge",
2937 tgt.value(i)
2938 );
2939 assert!((weight.value(i) - 1.0).abs() < 1e-9);
2940 }
2941 }
2942
2943 assert!(n_exact + n_fuzzy > 0, "no dup-cluster edges emitted");
2948 assert!(output.stats.exact_dups + output.stats.fuzzy_dups > 0);
2949
2950 let _ = std::fs::remove_dir_all(&dir);
2951 }
2952
2953 #[test]
2954 fn test_run_pipeline_deterministic_stats() {
2955 let dir_a = scratch_dir("det_a");
2956 let dir_b = scratch_dir("det_b");
2957 let a = run_kyc(&dir_a, "parquet", false);
2958 let b = run_kyc(&dir_b, "parquet", false);
2959 assert_eq!(a.stats.total_records, b.stats.total_records);
2960 assert_eq!(a.stats.exact_dups, b.stats.exact_dups);
2961 assert_eq!(a.stats.fuzzy_dups, b.stats.fuzzy_dups);
2962 assert_eq!(a.stats.hard_negs, b.stats.hard_negs);
2963 assert_eq!(a.stats.uniques, b.stats.uniques);
2964 assert_eq!(a.stats.masters, b.stats.masters);
2965 let _ = std::fs::remove_dir_all(&dir_a);
2966 let _ = std::fs::remove_dir_all(&dir_b);
2967 }
2968
2969 #[test]
2970 fn test_distribute_by_weight_sums_to_total_and_is_proportional() {
2971 let weights = [500_000usize, 216_666];
2979 let counts = distribute_by_weight(283_334, &weights);
2980 assert_eq!(counts.len(), 2);
2981 assert_eq!(counts.iter().sum::<usize>(), 283_334);
2982 assert!(counts[0] > 0, "first batch got zero duplicates");
2986 assert!(counts[1] > 0, "last batch got zero duplicates");
2987 assert!(counts[0] > counts[1]);
2988
2989 assert_eq!(distribute_by_weight(0, &weights), vec![0, 0]);
2991 assert_eq!(distribute_by_weight(100, &[0, 0]), vec![0, 0]);
2992 }
2993
2994 #[test]
2995 fn test_entity_prefix_unique_per_index() {
2996 let a = entity_prefix(0);
2997 let b = entity_prefix(1);
2998 assert_ne!(a, b);
2999 assert_eq!(a, "E00000");
3000 assert_eq!(entity_prefix(42), "E00042");
3001 }
3002
3003 #[test]
3004 fn test_run_pipeline_hell_multi_pass_touches_multiple_columns() {
3005 let dir = scratch_dir("hell_multipass");
3014 let schema = load_schema("kyc", &manifest_dir().join("schemas")).expect("load kyc.json");
3015 let mut ctx = Context::new(
3016 "kyc",
3017 "en",
3018 &manifest_dir().join("assets/pools").to_string_lossy(),
3019 )
3020 .expect("load context");
3021 let config = build_pipeline_config(
3022 "kyc",
3023 5000,
3024 7,
3025 "hell",
3026 0.1,
3027 0.3,
3028 &schema,
3029 "pipeline_test_hell_multipass",
3030 "ipc",
3031 false,
3032 "ipc",
3033 None,
3034 )
3035 .expect("build config");
3036 ctx.enable_watermark(&config.domain, config.size, config.seed);
3037 let output = run_pipeline(&ctx, &config, &dir.to_string_lossy()).expect("run_pipeline");
3038
3039 let read_all = |path: &str| -> RecordBatch {
3040 let file = std::fs::File::open(path).unwrap();
3041 let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
3042 let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
3043 let schema = batches[0].schema();
3044 let refs: Vec<Vec<&dyn Array>> = (0..schema.fields().len())
3045 .map(|i| batches.iter().map(|b| b.column(i).as_ref()).collect())
3046 .collect();
3047 let cols: Vec<ArrayRef> = refs
3048 .iter()
3049 .map(|r| arrow::compute::concat(r).unwrap())
3050 .collect();
3051 RecordBatch::try_new(schema, cols).unwrap()
3052 };
3053 let ds = read_all(&output.output_files[0]);
3054 let gt = read_all(&output.gt_file);
3055
3056 let ds_rid = ds.column_by_name("record_id").unwrap().as_string::<i32>();
3057 let ds_mid = ds.column_by_name("master_id").unwrap().as_string::<i32>();
3058 let gt_rid = gt.column_by_name("record_id").unwrap().as_string::<i32>();
3059 let gt_mt_col =
3062 arrow::compute::cast(gt.column_by_name("match_type").unwrap(), &DataType::Utf8)
3063 .unwrap();
3064 let gt_mt = gt_mt_col.as_string::<i32>();
3065
3066 let mut match_type_by_rid: HashMap<&str, &str> = HashMap::new();
3067 for i in 0..gt.num_rows() {
3068 match_type_by_rid.insert(gt_rid.value(i), gt_mt.value(i));
3069 }
3070
3071 let mut master_row_by_mid: HashMap<&str, usize> = HashMap::new();
3078 for i in 0..ds.num_rows() {
3079 master_row_by_mid.entry(ds_mid.value(i)).or_insert(i);
3080 }
3081
3082 let compare_cols: Vec<usize> = (0..ds.num_columns())
3083 .filter(|&i| {
3084 let field = ds.schema().field(i).clone();
3085 matches!(field.data_type(), DataType::Utf8)
3086 && !matches!(
3087 field.name().as_str(),
3088 "record_id" | "domain" | "entity_type" | "master_id"
3089 )
3090 && !field.name().ends_with("_id")
3091 })
3092 .collect();
3093
3094 let mut max_changed_cols = 0usize;
3095 for i in 0..ds.num_rows() {
3096 if match_type_by_rid.get(ds_rid.value(i)) != Some(&"fuzzy_dup") {
3097 continue;
3098 }
3099 let Some(&master_i) = master_row_by_mid.get(ds_mid.value(i)) else {
3100 continue;
3101 };
3102 let mut changed = 0usize;
3103 for &col_i in &compare_cols {
3104 let col = ds.column(col_i).as_string::<i32>();
3105 let a = (!col.is_null(i)).then(|| col.value(i));
3106 let b = (!col.is_null(master_i)).then(|| col.value(master_i));
3107 if a != b {
3108 changed += 1;
3109 }
3110 }
3111 max_changed_cols = max_changed_cols.max(changed);
3112 }
3113
3114 assert!(
3115 max_changed_cols >= 2,
3116 "expected at least one fuzzy_dup row to differ from its master on \
3117 2+ columns (evidence of multiple noise passes), got max {max_changed_cols}"
3118 );
3119
3120 let _ = std::fs::remove_dir_all(&dir);
3121 }
3122
3123 #[test]
3124 fn test_record_id_and_pad_string_roundtrip() {
3125 let rid = record_id_string(7);
3126 assert_eq!(rid, "R-0000000000007");
3127 let pad = pad_string(7);
3128 assert_eq!(pad, "0000000000007");
3129 let strs = record_id_strs(0..3);
3130 assert_eq!(strs.len(), 3);
3131 assert_eq!(strs[0], record_id_string(0));
3132 assert_eq!(strs[2], record_id_string(2));
3133 }
3134
3135 #[test]
3136 fn test_parse_record_idx_roundtrip() {
3137 for i in [0usize, 7, 1_234_567, 300_000_000] {
3138 assert_eq!(parse_record_idx(&record_id_string(i)), Some(i as u64));
3139 }
3140 assert_eq!(parse_record_idx("garbage"), None);
3141 assert_eq!(parse_record_idx("R-1"), None);
3142 assert_eq!(parse_record_idx("X-0000000000007"), None);
3143 }
3144
3145 #[test]
3146 fn test_pack_master_key_roundtrip_and_ordering() {
3147 let mid = |e: usize, i: usize| format!("{}-{}", entity_prefix(e), pad_string(i));
3148 assert!(pack_master_key(&mid(0, 0)).is_some());
3149 assert!(pack_master_key(&mid(0, 1)).unwrap() > pack_master_key(&mid(0, 0)).unwrap());
3151 assert!(
3154 pack_master_key(&mid(1, 0)).unwrap()
3155 > pack_master_key(&mid(0, 9_999_999_999_999)).unwrap()
3156 );
3157 assert_eq!(pack_master_key("garbage"), None);
3158 assert_eq!(pack_master_key("HN-0000000000007"), None);
3159 assert_eq!(pack_master_key("CANARY-000-0-abc"), None);
3160 }
3161
3162 #[test]
3172 #[ignore]
3173 fn hunt_h1_zip_cost_vs_column_count() {
3174 use arrow::array::StringArray;
3175 use arrow::datatypes::{DataType, Field, Schema};
3176
3177 fn make_batch(n_rows: usize, n_cols: usize) -> RecordBatch {
3178 let fields: Vec<Field> = (0..n_cols)
3179 .map(|i| Field::new(format!("col{i}"), DataType::Utf8, false))
3180 .collect();
3181 let cols: Vec<ArrayRef> = (0..n_cols)
3182 .map(|_| {
3183 Arc::new(StringArray::from(
3184 (0..n_rows)
3185 .map(|_| "abcdefghijklmnop".to_string())
3186 .collect::<Vec<_>>(),
3187 )) as ArrayRef
3188 })
3189 .collect();
3190 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
3191 }
3192
3193 const N_ROWS: usize = 100_000;
3194 const N_REPS: usize = 5;
3195 const EXTRA_PASSES: usize = 2;
3199
3200 eprintln!("\n[hunt_h1] N_ROWS={N_ROWS} N_REPS={N_REPS} EXTRA_PASSES={EXTRA_PASSES}");
3201 eprintln!("[hunt_h1] C=total columns, T=1 noised column (fixed) -- H1 predicts cost ∝ C");
3202
3203 for &n_cols in &[8usize, 15, 23] {
3204 let plan_cols = vec!["col0".to_string()];
3205 let all_types: Vec<&str> = vec!["typo"];
3206 let mut total = std::time::Duration::ZERO;
3207
3208 for rep in 0..N_REPS {
3209 let batch = make_batch(N_ROWS, n_cols);
3210 let mut rng = Rng::new(42 + rep as u64);
3211 let t0 = std::time::Instant::now();
3212
3213 let (noisy, mut unchanged, _) =
3214 apply_noise_with_retry(&batch, "typo", &plan_cols, &mut rng, &[]).unwrap();
3215 let mut current = noisy;
3216 for _ in 0..EXTRA_PASSES {
3217 current = apply_extra_pass_per_row(
3218 ¤t,
3219 &mut unchanged,
3220 &all_types,
3221 &mut rng,
3222 &[],
3223 )
3224 .unwrap();
3225 }
3226 std::hint::black_box(¤t);
3227 total += t0.elapsed();
3228 }
3229
3230 let avg = total / N_REPS as u32;
3231 eprintln!(
3232 "[hunt_h1] C={n_cols:2} avg={:>8.3}ms per_column={:>7.4}ms per_row_per_column={:>10.6}us",
3233 avg.as_secs_f64() * 1000.0,
3234 avg.as_secs_f64() * 1000.0 / n_cols as f64,
3235 avg.as_secs_f64() * 1_000_000.0 / (N_ROWS * n_cols) as f64,
3236 );
3237 }
3238 }
3239
3240 #[test]
3257 #[ignore]
3258 fn hunt_h1_residual_kyc_vs_nonprofit() {
3259 use arrow::array::StringArray;
3260 use arrow::datatypes::{DataType, Field, Schema};
3261
3262 fn make_batch(n_rows: usize, col_names: &[&str]) -> RecordBatch {
3263 let fields: Vec<Field> = col_names
3264 .iter()
3265 .map(|n| Field::new(*n, DataType::Utf8, true))
3266 .collect();
3267 let cols: Vec<ArrayRef> = col_names
3268 .iter()
3269 .map(|_| {
3270 Arc::new(StringArray::from(
3271 (0..n_rows)
3272 .map(|_| "abcdefghijklmnop testword".to_string())
3273 .collect::<Vec<_>>(),
3274 )) as ArrayRef
3275 })
3276 .collect();
3277 RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
3278 }
3279
3280 const N_ROWS: usize = 100_000;
3281 const N_REPS: usize = 5;
3282 const HELL_TYPES: &[&str] = &[
3283 "typo_aggressive",
3284 "visual",
3285 "unicode_pollution",
3286 "names",
3287 "dates",
3288 "identifiers",
3289 "addresses",
3290 "companies",
3291 "extra",
3292 ];
3293 const HELL_PASSES: usize = 8; let kyc_cols: &[&str] = &[
3296 "party_id",
3297 "given_name",
3298 "family_name",
3299 "middle_name",
3300 "birth_date",
3301 "birth_place",
3302 "nationality",
3303 "email_address",
3304 "mobile_phone",
3305 "personal_administrative_number",
3306 "tax_id",
3307 "residential_address_line1",
3308 "residential_city",
3309 "residential_postal_code",
3310 "residential_state",
3311 "residential_country",
3312 "occupation",
3313 "risk_score",
3314 "document_type",
3315 "document_number",
3316 "source_system",
3317 ];
3318 let nonprofit_cols: &[&str] = &[
3319 "donation_id",
3320 "donor_id",
3321 "campaign_id",
3322 "donation_date",
3323 "payment_method",
3324 "source_system",
3325 ];
3326
3327 eprintln!("\n[hunt_h1_residual] N_ROWS={N_ROWS} N_REPS={N_REPS} HELL_PASSES={HELL_PASSES}");
3328
3329 {
3334 let batch = make_batch(N_ROWS, kyc_cols);
3335 eprintln!("[hunt_h1_residual] per-category breakdown on kyc.natural_person:");
3336 for &ttype in HELL_TYPES {
3337 let mut rng = Rng::new(42);
3338 let t0 = std::time::Instant::now();
3339 let (_, _, tcols) =
3340 apply_noise_with_retry(&batch, ttype, &[], &mut rng, &[]).unwrap();
3341 eprintln!(
3342 "[hunt_h1_residual] {ttype:20} {:>8.3}ms target_cols={:?}",
3343 t0.elapsed().as_secs_f64() * 1000.0,
3344 tcols,
3345 );
3346 }
3347 }
3348
3349 for (label, col_names) in [
3350 ("kyc.natural_person", kyc_cols),
3351 ("nonprofit.donation", nonprofit_cols),
3352 ] {
3353 let mut total = std::time::Duration::ZERO;
3354 let mut total_target_hits = 0usize;
3355
3356 for rep in 0..N_REPS {
3357 let batch = make_batch(N_ROWS, col_names);
3358 let mut rng = Rng::new(42 + rep as u64);
3359 let t0 = std::time::Instant::now();
3360
3361 let (noisy, mut unchanged, t0cols) =
3364 apply_noise_with_retry(&batch, "identifiers", &[], &mut rng, &[]).unwrap();
3365 total_target_hits += t0cols.len();
3366 let mut current = noisy;
3367 for _ in 1..HELL_PASSES {
3368 current = apply_extra_pass_per_row(
3369 ¤t,
3370 &mut unchanged,
3371 HELL_TYPES,
3372 &mut rng,
3373 &[],
3374 )
3375 .unwrap();
3376 }
3377 std::hint::black_box(¤t);
3378 total += t0.elapsed();
3379 }
3380
3381 let avg = total / N_REPS as u32;
3382 eprintln!(
3383 "[hunt_h1_residual] {label:20} C={:2} avg={:>8.3}ms avg_first_pass_target_cols={:.1}",
3384 col_names.len(),
3385 avg.as_secs_f64() * 1000.0,
3386 total_target_hits as f64 / N_REPS as f64,
3387 );
3388 }
3389 }
3390}