Skip to main content

dupehell_core/
gt.rs

1// DupeHell -- MIT License . Educational Use Only
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// EDUCATIONAL AND RESEARCH PURPOSES ONLY -- see ETHICS.md for prohibited uses.
5// No liability for misuse.
6
7use 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
13/// Result of [`GtAccumulator::finish`].
14pub struct GtResult {
15    /// Cluster members that are genuinely byte-for-byte identical to their
16    /// master (the master itself, plus any duplicate copy whose assigned
17    /// noise ended up a no-op — see `pipeline::unchanged_row_mask`).
18    pub n_exact_dup: usize,
19    /// Cluster members that are duplicates of a master but differ from it
20    /// on at least one column (noise was actually applied). Distinguished
21    /// from `n_exact_dup` so a consumer can't mistake a fuzzy duplicate for
22    /// a trivial byte-for-byte match.
23    pub n_fuzzy_dup: usize,
24    pub n_hard_neg: usize,
25    pub n_unique: usize,
26    pub n_masters: usize,
27    /// Maps each duplicated `master_id` to the full set of `(record_id,
28    /// is_identical)` pairs in its cluster (base + duplicate copies, exact
29    /// and fuzzy alike) — `is_identical` mirrors that record's own
30    /// `exact_dup`/`fuzzy_dup` classification. Consumed by
31    /// `graph_gen::push_dup_clusters` to decide, per edge, whether the pair
32    /// it connects is `exact_dup` (both ends byte-identical to the master,
33    /// hence to each other) or `fuzzy_dup` (at least one end was noised).
34    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        // `true` for duplicate-copy rows whose assigned noise ended up not
43        // changing anything visible; `false` for duplicate copies that were
44        // genuinely altered, and for base/master rows (a base row is never
45        // itself "a copy that happened to match" — see `is_base` below for
46        // how its match_type is actually resolved). Meaningless for
47        // hard_neg/canary rows (classified by master_id prefix regardless).
48        Field::new("is_identical", DataType::Boolean, false),
49        // `true` only for base/master rows (one per entity, written by
50        // `push_base_batch`); `false` for every duplicate copy and for
51        // hard_neg/canary rows. Lets `finish` tell a base row apart from a
52        // genuine fuzzy copy when both carry `is_identical = false` — without
53        // it, a duplicated master's own base row would be misclassified
54        // `exact_dup` unconditionally (it used to hardcode
55        // `is_identical = true` for every base row, regardless of whether
56        // any of its actual duplicate copies were byte-identical to it).
57        Field::new("is_base", DataType::Boolean, false),
58    ]))
59}
60
61/// Streaming, two-pass ground-truth builder.
62///
63/// Duplicate detection needs a genuinely global view of which masters are
64/// duplicated (a master can be duplicated anywhere in the record stream),
65/// but that's the *only* thing that needs to be global — everything else
66/// can be processed one batch at a time. Previously the whole run
67/// accumulated three full-dataset `Vec<ArrayRef>` (record_id/entity_type/
68/// master_id) in RAM to concat and classify at the very end; at 200M+
69/// records that's multiple GB held for the entire run just for GT
70/// bookkeeping.
71///
72/// This accumulator instead:
73/// 1. streams each batch straight to a small "draft" IPC file as it's
74///    generated. The pipeline already knows, at the point each batch is
75///    produced, whether its rows introduce brand-new masters
76///    (`push_base_batch`) or are duplicate copies of masters seen earlier
77///    (`push_dup_batch`) — that's how `pipeline::run_pipeline` builds
78///    `dup_mids_buf` today. So instead of *discovering* duplication by
79///    folding occurrence counts over every master ever seen (O(n_base) —
80///    one entry per base record, duplicated or not), this only records the
81///    (much smaller) set of masters that are *actually* duplicated —
82///    O(n_dup_masters), a fraction of the dataset determined by the
83///    doublet/triplet ratios, not the full base population. HN/CANARY rows
84///    (`push_other_batch`) need no bookkeeping at all: their prefix alone
85///    determines their classification.
86/// 2. on `finish`, re-reads that draft file batch by batch — now that the
87///    full duplicated-master set is known — to classify each row and
88///    stream the definitive ground-truth file (IPC or Parquet), same
89///    batch-by-batch pattern already used elsewhere in this codebase for
90///    the IPC→Parquet dataset conversion (see `pipeline::run_pipeline`).
91///
92/// Duplicate detection is keyed on the **full** master_id string (entity
93/// prefix included), not just a numeric suffix: master_ids are assigned
94/// per-entity-plan starting from index 0, so two unrelated entities of
95/// different types can share the same numeric suffix. Keying on the full
96/// string avoids counting those as duplicates of each other.
97pub 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 that have at least one duplicate copy whose assigned noise
103    /// ended up a genuine no-op (byte-identical to the base) — as opposed to
104    /// `dup_masters`, which just means "has any duplicate copy at all,
105    /// exact or fuzzy". Used at `finish` so a base row can tell whether it
106    /// actually has an identical twin instead of unconditionally claiming
107    /// one.
108    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    /// Feed a batch of rows that each introduce a brand-new, not-yet-seen
154    /// master (base records — one master per row, never a duplicate copy).
155    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        // `is_identical = false`: a base row is not itself a "copy that
168        // happened to match" — whether it should read as `exact_dup` is
169        // resolved at `finish` from `masters_with_exact_copy`, not hardcoded
170        // here (this used to be `true` unconditionally, which mislabeled
171        // every duplicated master's base row `exact_dup` even when all of
172        // its actual duplicate copies were fuzzy — none byte-identical).
173        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    /// Feed a batch of duplicate-copy rows (each `master_id` matches an
179    /// existing base record already fed via `push_base_batch`). Records the
180    /// *set* of duplicated masters — every row sharing one of these
181    /// master_ids (the original base row included) belongs to a duplicated
182    /// cluster at `finish`, classified `exact_dup` or `fuzzy_dup` per-row
183    /// based on `is_identical` (`true` when the noise assigned to this copy
184    /// ended up a no-op — see `pipeline::unchanged_row_mask` — `false` when
185    /// it produced a real, visible change). Also records which masters have
186    /// at least one genuinely identical copy, for the base row's own
187    /// classification at `finish`.
188    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    /// Feed a batch of rows whose classification is fully determined by
222    /// their `master_id` prefix (hard negatives `HN-...`, canaries
223    /// `CANARY-...`) — no bookkeeping needed beyond streaming to the draft.
224    pub fn push_other_batch(
225        &mut self,
226        record_ids: &ArrayRef,
227        entity_types: &ArrayRef,
228        master_ids: &ArrayRef,
229    ) -> Result<(), String> {
230        // `is_identical`/`is_base` are meaningless for hard_neg/canary rows
231        // (classified by master_id prefix regardless), so the values don't
232        // matter.
233        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    /// Consumes the accumulator: closes the draft, re-reads it batch by
238    /// batch to classify each row now that the full duplicated-master set
239    /// is known, and streams the definitive ground-truth file (IPC or
240    /// Parquet). Returns a [`GtResult`].
241    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        // Cluster membership for duplicated masters (base + duplicate copies,
284        // exact and fuzzy alike), for emitting duplicate-cluster edges after
285        // this pass.
286        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                // Per-row, not per-cluster: a cluster can mix a base row, a
305                // copy the noise happened not to change, and a copy that's
306                // genuinely different. A base row has no noise of its own —
307                // it reads `exact_dup` only if the cluster actually contains
308                // a byte-identical copy (`masters_with_exact_copy`), never
309                // unconditionally (that used to mislabel every duplicated
310                // master's base row `exact_dup` even when all of its copies
311                // were fuzzy).
312                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                // Every row of a duplicated master belongs to its cluster,
336                // tagged with its own identical/fuzzy status.
337                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    /// Reads the final GT file into a `record_id -> match_type` map. A map
458    /// (rather than a positional `Vec`) keeps the assertions independent of
459    /// draft row order, which is caller-determined (base/dup/other batches
460    /// are pushed as separate, non-interleaved groups in the real pipeline).
461    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        // Base rows: R1 (master M-0000001, later duplicated), R3 (singleton
489        // M-0000002), R6 (singleton M-0000005).
490        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        // Dup row: R2 duplicates M-0000001, unchanged by its noise pass.
497        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        // Hard negatives: R4, R5.
505        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); // M-0000001, M-0000002, M-0000005
525
526        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    /// A duplicate copy whose noise pass actually changed something
538    /// (`is_identical = false`) must be classified `fuzzy_dup`, not
539    /// `exact_dup` — the master itself, and any sibling copy that was left
540    /// unchanged, keep `exact_dup` independently, per row.
541    #[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        // Triplet: R1 is the master, R2 was left unchanged by its noise
548        // pass, R3 was genuinely altered.
549        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); // R1 (master) + R2 (unchanged copy)
573        assert_eq!(fd, 1); // R3 (genuinely noised copy)
574        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        // Base batch: two masters, one of which (M-0000001) is duplicated
592        // in a later, separate dup batch — simulating a master's duplicate
593        // landing in a different batch than its base row.
594        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); // both rows of M-0000001
618        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    /// Regression: two different entity types can produce master_ids that
626    /// share the same numeric suffix (each entity's index restarts at 0),
627    /// but must not be counted as duplicates of each other.
628    #[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    /// `finish()` returns a `cluster_map` mapping each duplicated master_id
657    /// to the full set of `(record_id, is_identical)` pairs in its cluster
658    /// (base + duplicate copies). This is the structure consumed by
659    /// `push_dup_clusters` to decide, per edge, `exact_dup` vs `fuzzy_dup`.
660    #[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        // R2 (M-0000001) stayed identical; R7 (M-0000005) was genuinely noised.
673        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        // Only duplicated masters appear; singletons (M-0000002) do not.
694        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        // R6 is the base of M-0000005; its only copy (R7) was genuinely
701        // noised, so the cluster has no byte-identical pair at all — R6
702        // must NOT read as identical just because it's the base (that was
703        // the bug: base rows used to hardcode `is_identical = true`).
704        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}