1use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, StringArray, StringBuilder};
8use arrow::datatypes::{DataType, Field, Schema};
9use arrow::record_batch::RecordBatch;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13pub struct GtResult {
15 pub n_exact_dup: usize,
19 pub n_fuzzy_dup: usize,
24 pub n_hard_neg: usize,
25 pub n_unique: usize,
26 pub n_masters: usize,
27 pub cluster_map: HashMap<String, Vec<(String, bool)>>,
35}
36
37fn draft_schema() -> Arc<Schema> {
38 Arc::new(Schema::new(vec![
39 Field::new("record_id", DataType::Utf8, false),
40 Field::new("master_id", DataType::Utf8, false),
41 Field::new("entity_type", DataType::Utf8, false),
42 Field::new("is_identical", DataType::Boolean, false),
49 Field::new("is_base", DataType::Boolean, false),
58 ]))
59}
60
61pub struct GtAccumulator {
98 draft_path: String,
99 writer: arrow::ipc::writer::FileWriter<std::fs::File>,
100 schema: Arc<Schema>,
101 dup_masters: rustc_hash::FxHashSet<String>,
102 masters_with_exact_copy: rustc_hash::FxHashSet<String>,
109 n_base_masters: usize,
110}
111
112impl GtAccumulator {
113 pub fn new(draft_path: &str) -> Result<Self, String> {
114 let schema = draft_schema();
115 let file = std::fs::File::create(draft_path)
116 .map_err(|e| format!("create gt draft {draft_path}: {e}"))?;
117 let writer = arrow::ipc::writer::FileWriter::try_new(file, &schema)
118 .map_err(|e| format!("gt draft writer: {e}"))?;
119 Ok(Self {
120 draft_path: draft_path.to_string(),
121 writer,
122 schema,
123 dup_masters: rustc_hash::FxHashSet::default(),
124 masters_with_exact_copy: rustc_hash::FxHashSet::default(),
125 n_base_masters: 0,
126 })
127 }
128
129 fn write_draft(
130 &mut self,
131 record_ids: &ArrayRef,
132 entity_types: &ArrayRef,
133 master_ids: &ArrayRef,
134 is_identical: &ArrayRef,
135 is_base: &ArrayRef,
136 ) -> Result<(), String> {
137 let batch = RecordBatch::try_new(
138 self.schema.clone(),
139 vec![
140 record_ids.clone(),
141 master_ids.clone(),
142 entity_types.clone(),
143 is_identical.clone(),
144 is_base.clone(),
145 ],
146 )
147 .map_err(|e| format!("build gt draft batch: {e}"))?;
148 self.writer
149 .write(&batch)
150 .map_err(|e| format!("write gt draft: {e}"))
151 }
152
153 pub fn push_base_batch(
156 &mut self,
157 record_ids: &ArrayRef,
158 entity_types: &ArrayRef,
159 master_ids: &ArrayRef,
160 ) -> Result<(), String> {
161 let mids = master_ids.as_string::<i32>();
162 for i in 0..mids.len() {
163 if !mids.is_null(i) {
164 self.n_base_masters += 1;
165 }
166 }
167 let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; mids.len()]));
174 let all_true: ArrayRef = Arc::new(BooleanArray::from(vec![true; mids.len()]));
175 self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_true)
176 }
177
178 pub fn push_dup_batch(
189 &mut self,
190 record_ids: &ArrayRef,
191 entity_types: &ArrayRef,
192 master_ids: &ArrayRef,
193 is_identical: &ArrayRef,
194 ) -> Result<(), String> {
195 let mids = master_ids.as_string::<i32>();
196 let idents = is_identical.as_boolean();
197 for i in 0..mids.len() {
198 if !mids.is_null(i) {
199 let mid = mids.value(i);
200 if !self.dup_masters.contains(mid) {
201 self.dup_masters.insert(mid.to_string());
202 }
203 if !idents.is_null(i)
204 && idents.value(i)
205 && !self.masters_with_exact_copy.contains(mid)
206 {
207 self.masters_with_exact_copy.insert(mid.to_string());
208 }
209 }
210 }
211 let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; mids.len()]));
212 self.write_draft(
213 record_ids,
214 entity_types,
215 master_ids,
216 is_identical,
217 &all_false,
218 )
219 }
220
221 pub fn push_other_batch(
225 &mut self,
226 record_ids: &ArrayRef,
227 entity_types: &ArrayRef,
228 master_ids: &ArrayRef,
229 ) -> Result<(), String> {
230 let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; record_ids.len()]));
234 self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_false)
235 }
236
237 pub fn finish(
242 self,
243 difficulty: &str,
244 output_format: &str,
245 final_path: &str,
246 metadata: &HashMap<String, String>,
247 ) -> Result<GtResult, String> {
248 let GtAccumulator {
249 draft_path,
250 mut writer,
251 dup_masters,
252 masters_with_exact_copy,
253 n_base_masters,
254 ..
255 } = self;
256 writer
257 .finish()
258 .map_err(|e| format!("finish gt draft: {e}"))?;
259 drop(writer);
260
261 let final_schema = Arc::new(
262 Schema::new(vec![
263 Field::new("record_id", DataType::Utf8, false),
264 Field::new("master_id", DataType::Utf8, false),
265 Field::new("entity_type", DataType::Utf8, false),
266 Field::new("match_type", DataType::Utf8, false),
267 Field::new("difficulty", DataType::Utf8, false),
268 ])
269 .with_metadata(metadata.clone()),
270 );
271
272 let draft_file = std::fs::File::open(&draft_path)
273 .map_err(|e| format!("reopen gt draft {draft_path}: {e}"))?;
274 let reader = arrow::ipc::reader::FileReader::try_new(draft_file, None)
275 .map_err(|e| format!("gt draft reader: {e}"))?;
276
277 let mut sink = GtSink::new(output_format, final_path, &final_schema, metadata)?;
278
279 let mut n_exact_dup = 0usize;
280 let mut n_fuzzy_dup = 0usize;
281 let mut n_hard_neg = 0usize;
282 let mut n_unique = 0usize;
283 let mut cluster_map: HashMap<String, Vec<(String, bool)>> = HashMap::new();
287
288 for batch_result in reader {
289 let batch = batch_result.map_err(|e| format!("read gt draft batch: {e}"))?;
290 let n = batch.num_rows();
291 let mid_col = batch.column(1).as_string::<i32>();
292 let rid_col = batch.column(0).as_string::<i32>();
293 let ident_col = batch.column(3).as_boolean();
294 let base_col = batch.column(4).as_boolean();
295
296 let mut mt_builder = StringBuilder::with_capacity(n, n * 10);
297 for i in 0..n {
298 let mid = if mid_col.is_null(i) {
299 ""
300 } else {
301 mid_col.value(i)
302 };
303 let is_base = !base_col.is_null(i) && base_col.value(i);
304 let is_identical = if is_base {
313 masters_with_exact_copy.contains(mid)
314 } else {
315 !ident_col.is_null(i) && ident_col.value(i)
316 };
317 let mt = if mid.starts_with("HN-") {
318 n_hard_neg += 1;
319 "hard_neg"
320 } else if mid.starts_with("CANARY-") {
321 "canary"
322 } else if dup_masters.contains(mid) {
323 if is_identical {
324 n_exact_dup += 1;
325 "exact_dup"
326 } else {
327 n_fuzzy_dup += 1;
328 "fuzzy_dup"
329 }
330 } else {
331 n_unique += 1;
332 "unique"
333 };
334 mt_builder.append_value(mt);
335 if dup_masters.contains(mid) {
338 cluster_map
339 .entry(mid.to_string())
340 .or_default()
341 .push((rid_col.value(i).to_string(), is_identical));
342 }
343 }
344 let mt_arr: ArrayRef = Arc::new(mt_builder.finish());
345 let diff_arr: ArrayRef = Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
346 difficulty, n,
347 )));
348
349 let final_batch = RecordBatch::try_new(
350 final_schema.clone(),
351 vec![
352 batch.column(0).clone(),
353 batch.column(1).clone(),
354 batch.column(2).clone(),
355 mt_arr,
356 diff_arr,
357 ],
358 )
359 .map_err(|e| format!("build final gt batch: {e}"))?;
360
361 sink.write(&final_batch)?;
362 }
363
364 sink.finish()?;
365 std::fs::remove_file(&draft_path).ok();
366
367 Ok(GtResult {
368 n_exact_dup,
369 n_fuzzy_dup,
370 n_hard_neg,
371 n_unique,
372 n_masters: n_base_masters,
373 cluster_map,
374 })
375 }
376}
377
378enum GtSink {
379 Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
380 Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
381}
382
383impl GtSink {
384 fn new(
385 output_format: &str,
386 path: &str,
387 schema: &Arc<Schema>,
388 metadata: &HashMap<String, String>,
389 ) -> Result<Self, String> {
390 let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
391 if output_format == "parquet" {
392 use parquet::basic::{Compression, ZstdLevel};
393 use parquet::file::properties::WriterProperties;
394 let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
395 let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
396 .iter()
397 .map(|(k, v)| parquet::file::metadata::KeyValue {
398 key: k.clone(),
399 value: Some(v.clone()),
400 })
401 .collect();
402 let props = WriterProperties::builder()
403 .set_compression(Compression::ZSTD(zstd))
404 .set_data_page_size_limit(1_048_576)
405 .set_key_value_metadata(Some(meta_kv))
406 .build();
407 let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
408 .map_err(|e| format!("gt parquet writer: {e}"))?;
409 Ok(GtSink::Parquet(Box::new(writer)))
410 } else {
411 let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
412 .map_err(|e| format!("gt ipc writer: {e}"))?;
413 Ok(GtSink::Ipc(Box::new(writer)))
414 }
415 }
416
417 fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
418 match self {
419 GtSink::Ipc(w) => w.write(batch).map_err(|e| format!("write gt ipc: {e}")),
420 GtSink::Parquet(w) => w.write(batch).map_err(|e| format!("write gt parquet: {e}")),
421 }
422 }
423
424 fn finish(self) -> Result<(), String> {
425 match self {
426 GtSink::Ipc(mut w) => w.finish().map_err(|e| format!("finish gt ipc: {e}")),
427 GtSink::Parquet(w) => w
428 .close()
429 .map(|_| ())
430 .map_err(|e| format!("close gt parquet: {e}")),
431 }
432 }
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438 use arrow::array::StringArray;
439
440 fn tmp_path(name: &str) -> String {
441 let mut p = std::env::temp_dir();
442 p.push(format!(
443 "dupehell_gt_test_{name}_{}.ipc",
444 std::process::id()
445 ));
446 p.to_string_lossy().into_owned()
447 }
448
449 fn arr(values: Vec<&str>) -> ArrayRef {
450 Arc::new(StringArray::from(values))
451 }
452
453 fn barr(values: Vec<bool>) -> ArrayRef {
454 Arc::new(BooleanArray::from(values))
455 }
456
457 fn read_match_types(final_path: &str) -> HashMap<String, String> {
462 let file = std::fs::File::open(final_path).unwrap();
463 let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
464 let mut out = HashMap::new();
465 for batch in reader {
466 let batch = batch.unwrap();
467 let rid = batch
468 .column_by_name("record_id")
469 .unwrap()
470 .as_string::<i32>();
471 let mt = batch
472 .column_by_name("match_type")
473 .unwrap()
474 .as_string::<i32>();
475 for i in 0..batch.num_rows() {
476 out.insert(rid.value(i).to_string(), mt.value(i).to_string());
477 }
478 }
479 out
480 }
481
482 #[test]
483 fn test_gt_accumulator_basic() {
484 let draft = tmp_path("basic_draft");
485 let final_path = tmp_path("basic_final");
486
487 let mut acc = GtAccumulator::new(&draft).unwrap();
488 acc.push_base_batch(
491 &arr(vec!["R1", "R3", "R6"]),
492 &arr(vec!["person", "person", "person"]),
493 &arr(vec!["M-0000001", "M-0000002", "M-0000005"]),
494 )
495 .unwrap();
496 acc.push_dup_batch(
498 &arr(vec!["R2"]),
499 &arr(vec!["person"]),
500 &arr(vec!["M-0000001"]),
501 &barr(vec![true]),
502 )
503 .unwrap();
504 acc.push_other_batch(
506 &arr(vec!["R4", "R5"]),
507 &arr(vec!["person", "person"]),
508 &arr(vec!["HN-0000003", "HN-0000004"]),
509 )
510 .unwrap();
511
512 let GtResult {
513 n_exact_dup: ed,
514 n_hard_neg: hn,
515 n_unique: un,
516 n_masters: masters,
517 ..
518 } = acc
519 .finish("medium", "ipc", &final_path, &HashMap::new())
520 .unwrap();
521 assert_eq!(ed, 2);
522 assert_eq!(hn, 2);
523 assert_eq!(un, 2);
524 assert_eq!(masters, 3); let match_types = read_match_types(&final_path);
527 assert_eq!(match_types["R1"], "exact_dup");
528 assert_eq!(match_types["R2"], "exact_dup");
529 assert_eq!(match_types["R3"], "unique");
530 assert_eq!(match_types["R4"], "hard_neg");
531 assert_eq!(match_types["R5"], "hard_neg");
532 assert_eq!(match_types["R6"], "unique");
533
534 std::fs::remove_file(&final_path).ok();
535 }
536
537 #[test]
542 fn test_gt_accumulator_fuzzy_dup() {
543 let draft = tmp_path("fuzzy_draft");
544 let final_path = tmp_path("fuzzy_final");
545
546 let mut acc = GtAccumulator::new(&draft).unwrap();
547 acc.push_base_batch(
550 &arr(vec!["R1"]),
551 &arr(vec!["person"]),
552 &arr(vec!["M-0000001"]),
553 )
554 .unwrap();
555 acc.push_dup_batch(
556 &arr(vec!["R2", "R3"]),
557 &arr(vec!["person", "person"]),
558 &arr(vec!["M-0000001", "M-0000001"]),
559 &barr(vec![true, false]),
560 )
561 .unwrap();
562
563 let GtResult {
564 n_exact_dup: ed,
565 n_fuzzy_dup: fd,
566 n_unique: un,
567 n_masters: masters,
568 ..
569 } = acc
570 .finish("medium", "ipc", &final_path, &HashMap::new())
571 .unwrap();
572 assert_eq!(ed, 2); assert_eq!(fd, 1); assert_eq!(un, 0);
575 assert_eq!(masters, 1);
576
577 let match_types = read_match_types(&final_path);
578 assert_eq!(match_types["R1"], "exact_dup");
579 assert_eq!(match_types["R2"], "exact_dup");
580 assert_eq!(match_types["R3"], "fuzzy_dup");
581
582 std::fs::remove_file(&final_path).ok();
583 }
584
585 #[test]
586 fn test_gt_accumulator_multi_batch() {
587 let draft = tmp_path("multi_draft");
588 let final_path = tmp_path("multi_final");
589
590 let mut acc = GtAccumulator::new(&draft).unwrap();
591 acc.push_base_batch(
595 &arr(vec!["R1", "R3"]),
596 &arr(vec!["person", "person"]),
597 &arr(vec!["M-0000001", "M-0000002"]),
598 )
599 .unwrap();
600 acc.push_dup_batch(
601 &arr(vec!["R2"]),
602 &arr(vec!["person"]),
603 &arr(vec!["M-0000001"]),
604 &barr(vec![true]),
605 )
606 .unwrap();
607
608 let GtResult {
609 n_exact_dup: ed,
610 n_hard_neg: hn,
611 n_unique: un,
612 n_masters: masters,
613 ..
614 } = acc
615 .finish("medium", "ipc", &final_path, &HashMap::new())
616 .unwrap();
617 assert_eq!(ed, 2); assert_eq!(hn, 0);
619 assert_eq!(un, 1);
620 assert_eq!(masters, 2);
621
622 std::fs::remove_file(&final_path).ok();
623 }
624
625 #[test]
629 fn test_no_cross_entity_suffix_collision() {
630 let draft = tmp_path("suffix_draft");
631 let final_path = tmp_path("suffix_final");
632
633 let mut acc = GtAccumulator::new(&draft).unwrap();
634 acc.push_base_batch(
635 &arr(vec!["R1", "R2"]),
636 &arr(vec!["person", "account"]),
637 &arr(vec!["PERSON-0000001", "ACCOUNT-0000001"]),
638 )
639 .unwrap();
640
641 let GtResult {
642 n_exact_dup: ed,
643 n_unique: un,
644 n_masters: masters,
645 ..
646 } = acc
647 .finish("medium", "ipc", &final_path, &HashMap::new())
648 .unwrap();
649 assert_eq!(ed, 0);
650 assert_eq!(un, 2);
651 assert_eq!(masters, 2);
652
653 std::fs::remove_file(&final_path).ok();
654 }
655
656 #[test]
661 fn test_cluster_map_contents() {
662 let draft = tmp_path("cm_draft");
663 let final_path = tmp_path("cm_final");
664
665 let mut acc = GtAccumulator::new(&draft).unwrap();
666 acc.push_base_batch(
667 &arr(vec!["R1", "R3", "R6"]),
668 &arr(vec!["person", "person", "person"]),
669 &arr(vec!["M-0000001", "M-0000002", "M-0000005"]),
670 )
671 .unwrap();
672 acc.push_dup_batch(
674 &arr(vec!["R2", "R7"]),
675 &arr(vec!["person", "person"]),
676 &arr(vec!["M-0000001", "M-0000005"]),
677 &barr(vec![true, false]),
678 )
679 .unwrap();
680 acc.push_other_batch(
681 &arr(vec!["R4", "R5"]),
682 &arr(vec!["person", "person"]),
683 &arr(vec!["HN-0000003", "HN-0000004"]),
684 )
685 .unwrap();
686
687 let GtResult {
688 cluster_map: cm, ..
689 } = acc
690 .finish("medium", "ipc", &final_path, &HashMap::new())
691 .unwrap();
692
693 assert_eq!(cm.len(), 2);
695 let mut m1 = cm.get("M-0000001").unwrap().clone();
696 m1.sort();
697 assert_eq!(m1, vec![("R1".to_string(), true), ("R2".to_string(), true)]);
698 let mut m5 = cm.get("M-0000005").unwrap().clone();
699 m5.sort();
700 assert_eq!(
705 m5,
706 vec![("R6".to_string(), false), ("R7".to_string(), false)]
707 );
708 assert!(!cm.contains_key("M-0000002"));
709
710 std::fs::remove_file(&final_path).ok();
711 }
712}