dupehell_core/gt.rs
1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
7use arrow::buffer::BooleanBuffer;
8use arrow::datatypes::{DataType, Field, Schema};
9use arrow::record_batch::RecordBatch;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13/// Compressed-sparse-row representation of the duplicate-cluster
14/// membership previously held as `HashMap<String, Vec<(String, bool)>>`.
15/// Built once, in `GtAccumulator::finish`, from every `(master_id,
16/// record_id, is_identical)` triple belonging to a duplicated master.
17///
18/// Avoids one `String` allocation per `record_id` and one per distinct
19/// `master_id` (the dominant RAM cost of `cluster_map` at 200M+ records,
20/// see `project_csr_cluster_map_backlog`): `master_id`/`record_id` are
21/// packed into `u64`s (`pipeline::pack_master_key`/`parse_record_idx` —
22/// both are pure functions of the fixed-width ID shapes `pipeline.rs`
23/// already generates, so the packing/unpacking is lossless and exact) and
24/// stored in two flat, contiguous buffers instead of a hash map of
25/// per-cluster `Vec`s.
26///
27/// `master_id` itself is never needed downstream as a string — the only
28/// consumer (`graph_gen::push_dup_clusters`) only ever used it to get a
29/// deterministic cluster iteration order (`master_ids.sort()`), and
30/// numeric comparison of the packed key preserves that exact order (both
31/// segments of the original string are zero-padded to a fixed width, so
32/// lexicographic and numeric order agree).
33/// Bit-packed boolean buffer — 1 bit per entry instead of `Vec<bool>`'s 1
34/// byte, an 8x reduction on `ClusterCsr::is_identical`, which parallels
35/// `records` (potentially tens of millions of entries at 100M+ hell scale
36/// with `--graph`, a low singleton fraction). Found via a triage of
37/// reusable Rust perf patterns (hunt3007/H7, Linkars trick #22).
38/// A `BooleanArray` of `n` copies of `value`, built directly from a
39/// `BooleanBuffer` (one bitmap allocation, `n/8` bytes) instead of
40/// `BooleanArray::from(vec![value; n])` (an `n`-byte `Vec<bool>` allocated
41/// and memset, then re-packed into a bitmap by `From`).
42fn const_bool_array(n: usize, value: bool) -> ArrayRef {
43 let buf = if value {
44 BooleanBuffer::new_set(n)
45 } else {
46 BooleanBuffer::new_unset(n)
47 };
48 Arc::new(BooleanArray::new(buf, None))
49}
50
51struct Bitset {
52 words: Vec<u64>,
53 len: usize,
54}
55
56impl Bitset {
57 fn with_capacity(n: usize) -> Self {
58 Self {
59 words: vec![0u64; n.div_ceil(64).max(1)],
60 len: 0,
61 }
62 }
63
64 fn push(&mut self, val: bool) {
65 let word_idx = self.len / 64;
66 if word_idx >= self.words.len() {
67 self.words.push(0);
68 }
69 if val {
70 self.words[word_idx] |= 1u64 << (self.len % 64);
71 }
72 self.len += 1;
73 }
74
75 fn get(&self, i: usize) -> bool {
76 (self.words[i / 64] >> (i % 64)) & 1 != 0
77 }
78}
79
80/// View into a contiguous range of a [`Bitset`], returned per-cluster by
81/// [`ClusterCsr::groups`] — parallel to the `&[u64]` record-index slice for
82/// the same cluster.
83pub struct BitspanRef<'a> {
84 bits: &'a Bitset,
85 start: usize,
86 end: usize,
87}
88
89impl BitspanRef<'_> {
90 /// `i` is local to this span (`0..self.len()`), matching how callers
91 /// already index the parallel `&[u64]` records slice.
92 pub fn get(&self, i: usize) -> bool {
93 self.bits.get(self.start + i)
94 }
95
96 pub fn len(&self) -> usize {
97 self.end - self.start
98 }
99
100 pub fn is_empty(&self) -> bool {
101 self.start == self.end
102 }
103
104 #[cfg(test)]
105 fn to_vec(&self) -> Vec<bool> {
106 (0..self.len()).map(|i| self.get(i)).collect()
107 }
108}
109
110pub struct ClusterCsr {
111 /// `offsets.len() == n_clusters + 1`; cluster `k`'s members are
112 /// `records[offsets[k]..offsets[k+1]]` (and the parallel span of
113 /// `is_identical`). Sorted ascending by packed master key.
114 offsets: Vec<u32>,
115 /// Flat, cluster-grouped buffer of global record indices (see
116 /// `pipeline::record_id_string`/`parse_record_idx`), ascending within
117 /// each cluster.
118 records: Vec<u64>,
119 /// Parallel to `records`: whether that record is byte-identical to its
120 /// cluster's master.
121 is_identical: Bitset,
122}
123
124/// `record_idx` (from `parse_record_idx`) is bounded by `PAD_LEN` (13)
125/// decimal digits — `< 10^13 ≈ 2^43.2` — so bit 63 is always free to carry
126/// `is_identical` alongside it in one `u64`, the same trick
127/// `pack_master_key` already uses for its own high bits (hunt1808/H-RAM3).
128/// Packing this into `ClusterCsr::build`'s transient pair list drops it
129/// from `(u64, u64, bool)` (24 bytes/entry — `bool` pads the tuple to the
130/// next 8-byte alignment boundary, measured on this machine) to `(u64,
131/// u64)` (16 bytes/entry, −33%), at potentially tens of millions of
132/// entries (every row of every duplicated cluster, `--graph` mode).
133const RIDX_IDENTICAL_BIT: u64 = 1 << 63;
134const RIDX_MASK: u64 = !RIDX_IDENTICAL_BIT;
135
136/// Packs `(record_idx, is_identical)` into `ClusterCsr::build`'s pair
137/// representation — see `RIDX_IDENTICAL_BIT`.
138pub(crate) fn pack_ridx_identical(record_idx: u64, is_identical: bool) -> u64 {
139 debug_assert_eq!(
140 record_idx & RIDX_IDENTICAL_BIT,
141 0,
142 "record_idx must fit in the low 63 bits"
143 );
144 record_idx | if is_identical { RIDX_IDENTICAL_BIT } else { 0 }
145}
146
147impl ClusterCsr {
148 /// Builds the CSR from an unordered flat list of `(packed_master_key,
149 /// packed_ridx_identical)` pairs (see `pack_ridx_identical`). Sorting
150 /// once by `(master_key, record_idx)` — masking off the packed
151 /// `is_identical` bit for the comparison, so it can't perturb member
152 /// order within a cluster — both groups every pair by cluster and
153 /// orders each cluster's members ascending, in one pass; no need to
154 /// pre-count group sizes or build the structure incrementally.
155 ///
156 /// Sorted in parallel (`rayon`, already a warm dependency by this point
157 /// in the run): the sort key `(master_key, record_idx)` is unique per
158 /// entry (one record belongs to exactly one cluster, and a
159 /// `record_idx` is unique across the whole run), so an unstable
160 /// parallel sort produces the exact same total order as the previous
161 /// single-threaded one — determinism doesn't depend on stability here.
162 pub(crate) fn build(mut pairs: Vec<(u64, u64)>) -> Self {
163 use rayon::slice::ParallelSliceMut;
164 pairs.par_sort_unstable_by_key(|&(mk, r)| (mk, r & RIDX_MASK));
165
166 let mut offsets = Vec::new();
167 let mut records = Vec::with_capacity(pairs.len());
168 let mut is_identical = Bitset::with_capacity(pairs.len());
169 let mut last_key: Option<u64> = None;
170
171 for (mk, r) in pairs {
172 if last_key != Some(mk) {
173 offsets.push(records.len() as u32);
174 last_key = Some(mk);
175 }
176 records.push(r & RIDX_MASK);
177 is_identical.push(r & RIDX_IDENTICAL_BIT != 0);
178 }
179 offsets.push(records.len() as u32);
180
181 Self {
182 offsets,
183 records,
184 is_identical,
185 }
186 }
187
188 /// Iterates clusters in ascending packed-master-key order (matches the
189 /// previous `master_ids.sort()` determinism exactly), yielding each
190 /// cluster's `(record_indices, is_identical)` pair. A cluster's
191 /// `record_indices` are already ascending, so callers that used to sort
192 /// members themselves (`push_dup_clusters`) no longer need to.
193 pub fn groups(&self) -> impl Iterator<Item = (&[u64], BitspanRef<'_>)> {
194 self.offsets.windows(2).map(move |w| {
195 let (start, end) = (w[0] as usize, w[1] as usize);
196 (
197 &self.records[start..end],
198 BitspanRef {
199 bits: &self.is_identical,
200 start,
201 end,
202 },
203 )
204 })
205 }
206
207 #[cfg(test)]
208 fn n_clusters(&self) -> usize {
209 self.offsets.len().saturating_sub(1)
210 }
211}
212
213/// Result of [`GtAccumulator::finish`].
214pub struct GtResult {
215 /// Cluster members that are genuinely byte-for-byte identical to their
216 /// master (the master itself, plus any duplicate copy whose assigned
217 /// noise ended up a no-op — see `pipeline::unchanged_row_mask`).
218 pub n_exact_dup: usize,
219 /// Cluster members that are duplicates of a master but differ from it
220 /// on at least one column (noise was actually applied). Distinguished
221 /// from `n_exact_dup` so a consumer can't mistake a fuzzy duplicate for
222 /// a trivial byte-for-byte match.
223 pub n_fuzzy_dup: usize,
224 pub n_hard_neg: usize,
225 pub n_unique: usize,
226 pub n_masters: usize,
227 /// Maps each duplicated `master_id` to the full set of `(record_id,
228 /// is_identical)` pairs in its cluster (base + duplicate copies, exact
229 /// and fuzzy alike) — `is_identical` mirrors that record's own
230 /// `exact_dup`/`fuzzy_dup` classification. Consumed by
231 /// `graph_gen::push_dup_clusters` to decide, per edge, whether the pair
232 /// it connects is `exact_dup` (both ends byte-identical to the master,
233 /// hence to each other) or `fuzzy_dup` (at least one end was noised).
234 pub cluster_map: ClusterCsr,
235}
236
237/// `Dictionary(Int32, Utf8)` for `entity_type`/`match_type`/`difficulty` —
238/// hunt1808/H11, an explicit output-contract change (not bit-identical:
239/// the declared column type changes from `Utf8`), applied on explicit user
240/// request. `entity_type` here must match `pipeline::low_cardinality_dict_type`
241/// exactly — this column's data is `add_metadata_and_align`'s `et_arr`,
242/// pushed straight through by `push_*_batch`, so draft/final schema must
243/// declare the same type it's actually built as.
244fn low_cardinality_dict_type() -> DataType {
245 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
246}
247
248fn draft_schema() -> Arc<Schema> {
249 Arc::new(Schema::new(vec![
250 Field::new("record_id", DataType::Utf8, false),
251 Field::new("master_id", DataType::Utf8, false),
252 Field::new("entity_type", low_cardinality_dict_type(), false),
253 // `true` for duplicate-copy rows whose assigned noise ended up not
254 // changing anything visible; `false` for duplicate copies that were
255 // genuinely altered, and for base/master rows (a base row is never
256 // itself "a copy that happened to match" — see `is_base` below for
257 // how its match_type is actually resolved). Meaningless for
258 // hard_neg/canary rows (classified by master_id prefix regardless).
259 Field::new("is_identical", DataType::Boolean, false),
260 // `true` only for base/master rows (one per entity, written by
261 // `push_base_batch`); `false` for every duplicate copy and for
262 // hard_neg/canary rows. Lets `finish` tell a base row apart from a
263 // genuine fuzzy copy when both carry `is_identical = false` — without
264 // it, a duplicated master's own base row would be misclassified
265 // `exact_dup` unconditionally (it used to hardcode
266 // `is_identical = true` for every base row, regardless of whether
267 // any of its actual duplicate copies were byte-identical to it).
268 Field::new("is_base", DataType::Boolean, false),
269 ]))
270}
271
272/// Streaming, two-pass ground-truth builder.
273///
274/// Duplicate detection needs a genuinely global view of which masters are
275/// duplicated (a master can be duplicated anywhere in the record stream),
276/// but that's the *only* thing that needs to be global — everything else
277/// can be processed one batch at a time. Previously the whole run
278/// accumulated three full-dataset `Vec<ArrayRef>` (record_id/entity_type/
279/// master_id) in RAM to concat and classify at the very end; at 200M+
280/// records that's multiple GB held for the entire run just for GT
281/// bookkeeping.
282///
283/// This accumulator instead:
284/// 1. streams each batch straight to a small "draft" IPC file as it's
285/// generated. The pipeline already knows, at the point each batch is
286/// produced, whether its rows introduce brand-new masters
287/// (`push_base_batch`) or are duplicate copies of masters seen earlier
288/// (`push_dup_batch`) — that's how `pipeline::run_pipeline` builds
289/// `dup_mids_buf` today. So instead of *discovering* duplication by
290/// folding occurrence counts over every master ever seen (O(n_base) —
291/// one entry per base record, duplicated or not), this only records the
292/// (much smaller) set of masters that are *actually* duplicated —
293/// O(n_dup_masters), a fraction of the dataset determined by the
294/// doublet/triplet ratios, not the full base population. HN/CANARY rows
295/// (`push_other_batch`) need no bookkeeping at all: their prefix alone
296/// determines their classification.
297/// 2. on `finish`, re-reads that draft file batch by batch — now that the
298/// full duplicated-master set is known — to classify each row and
299/// stream the definitive ground-truth file (IPC or Parquet), same
300/// batch-by-batch pattern already used elsewhere in this codebase for
301/// the IPC→Parquet dataset conversion (see `pipeline::run_pipeline`).
302///
303/// Duplicate detection is keyed on the **full** master_id (entity prefix
304/// included, via `pipeline::pack_master_key` into a `u64`), not just a
305/// numeric suffix: master_ids are assigned per-entity-plan starting from
306/// index 0, so two unrelated entities of different types can share the
307/// same numeric suffix. Keying on the full identity avoids counting those
308/// as duplicates of each other.
309pub struct GtAccumulator {
310 draft_path: String,
311 writer: arrow::ipc::writer::FileWriter<std::fs::File>,
312 schema: Arc<Schema>,
313 /// Keyed on `pipeline::pack_master_key(mid)` rather than the raw
314 /// `String` — avoids a per-row string hash/allocation on a set queried
315 /// 2-3x per duplicated row (`push_dup_batch` + classification in
316 /// `finish`), which VTune hotspots (hunt3007) showed costing ~5% of
317 /// total CPU on its own. `push_dup_batch` never receives an
318 /// HN-/CANARY- master_id, so every key here is guaranteed
319 /// pack_master_key-compatible.
320 dup_masters: rustc_hash::FxHashSet<u64>,
321 /// Masters that have at least one duplicate copy whose assigned noise
322 /// ended up a genuine no-op (byte-identical to the base) — as opposed to
323 /// `dup_masters`, which just means "has any duplicate copy at all,
324 /// exact or fuzzy". Used at `finish` so a base row can tell whether it
325 /// actually has an identical twin instead of unconditionally claiming
326 /// one.
327 masters_with_exact_copy: rustc_hash::FxHashSet<u64>,
328 n_base_masters: usize,
329}
330
331impl GtAccumulator {
332 pub fn new(draft_path: &str) -> Result<Self, String> {
333 let schema = draft_schema();
334 let file = std::fs::File::create(draft_path)
335 .map_err(|e| format!("create gt draft {draft_path}: {e}"))?;
336 let writer = arrow::ipc::writer::FileWriter::try_new(file, &schema)
337 .map_err(|e| format!("gt draft writer: {e}"))?;
338 Ok(Self {
339 draft_path: draft_path.to_string(),
340 writer,
341 schema,
342 dup_masters: rustc_hash::FxHashSet::default(),
343 masters_with_exact_copy: rustc_hash::FxHashSet::default(),
344 n_base_masters: 0,
345 })
346 }
347
348 fn write_draft(
349 &mut self,
350 record_ids: &ArrayRef,
351 entity_types: &ArrayRef,
352 master_ids: &ArrayRef,
353 is_identical: &ArrayRef,
354 is_base: &ArrayRef,
355 ) -> Result<(), String> {
356 let batch = RecordBatch::try_new(
357 self.schema.clone(),
358 vec![
359 record_ids.clone(),
360 master_ids.clone(),
361 entity_types.clone(),
362 is_identical.clone(),
363 is_base.clone(),
364 ],
365 )
366 .map_err(|e| format!("build gt draft batch: {e}"))?;
367 self.writer
368 .write(&batch)
369 .map_err(|e| format!("write gt draft: {e}"))
370 }
371
372 /// Feed a batch of rows that each introduce a brand-new, not-yet-seen
373 /// master (base records — one master per row, never a duplicate copy).
374 pub fn push_base_batch(
375 &mut self,
376 record_ids: &ArrayRef,
377 entity_types: &ArrayRef,
378 master_ids: &ArrayRef,
379 ) -> Result<(), String> {
380 let mids = master_ids.as_string::<i32>();
381 self.n_base_masters += mids.len() - mids.null_count();
382 // `is_identical = false`: a base row is not itself a "copy that
383 // happened to match" — whether it should read as `exact_dup` is
384 // resolved at `finish` from `masters_with_exact_copy`, not hardcoded
385 // here (this used to be `true` unconditionally, which mislabeled
386 // every duplicated master's base row `exact_dup` even when all of
387 // its actual duplicate copies were fuzzy — none byte-identical).
388 let all_false = const_bool_array(mids.len(), false);
389 let all_true = const_bool_array(mids.len(), true);
390 self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_true)
391 }
392
393 /// Feed a batch of duplicate-copy rows (each `master_id` matches an
394 /// existing base record already fed via `push_base_batch`). Records the
395 /// *set* of duplicated masters — every row sharing one of these
396 /// master_ids (the original base row included) belongs to a duplicated
397 /// cluster at `finish`, classified `exact_dup` or `fuzzy_dup` per-row
398 /// based on `is_identical` (`true` when the noise assigned to this copy
399 /// ended up a no-op — see `pipeline::unchanged_row_mask` — `false` when
400 /// it produced a real, visible change). Also records which masters have
401 /// at least one genuinely identical copy, for the base row's own
402 /// classification at `finish`.
403 pub fn push_dup_batch(
404 &mut self,
405 record_ids: &ArrayRef,
406 entity_types: &ArrayRef,
407 master_ids: &ArrayRef,
408 is_identical: &ArrayRef,
409 ) -> Result<(), String> {
410 let mids = master_ids.as_string::<i32>();
411 let idents = is_identical.as_boolean();
412 for i in 0..mids.len() {
413 if !mids.is_null(i) {
414 let mid = mids.value(i);
415 let key = crate::pipeline::pack_master_key(mid).unwrap_or_else(|| {
416 panic!(
417 "push_dup_batch: master_id {mid:?} isn't the fixed \
418 \"{{entity_prefix}}-{{pad_string}}\" shape -- \
419 push_dup_batch never receives an HN-/CANARY- master_id"
420 )
421 });
422 self.dup_masters.insert(key);
423 if !idents.is_null(i) && idents.value(i) {
424 self.masters_with_exact_copy.insert(key);
425 }
426 }
427 }
428 let all_false = const_bool_array(mids.len(), false);
429 self.write_draft(
430 record_ids,
431 entity_types,
432 master_ids,
433 is_identical,
434 &all_false,
435 )
436 }
437
438 /// Feed a batch of rows whose classification is fully determined by
439 /// their `master_id` prefix (hard negatives `HN-...`, canaries
440 /// `CANARY-...`) — no bookkeeping needed beyond streaming to the draft.
441 pub fn push_other_batch(
442 &mut self,
443 record_ids: &ArrayRef,
444 entity_types: &ArrayRef,
445 master_ids: &ArrayRef,
446 ) -> Result<(), String> {
447 // `is_identical`/`is_base` are meaningless for hard_neg/canary rows
448 // (classified by master_id prefix regardless), so the values don't
449 // matter.
450 let all_false = const_bool_array(record_ids.len(), false);
451 self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_false)
452 }
453
454 /// Consumes the accumulator: closes the draft, re-reads it batch by
455 /// batch to classify each row now that the full duplicated-master set
456 /// is known, and streams the definitive ground-truth file (IPC or
457 /// Parquet). Returns a [`GtResult`].
458 pub fn finish(
459 self,
460 difficulty: &str,
461 output_format: &str,
462 final_path: &str,
463 metadata: &HashMap<String, String>,
464 track_clusters: bool,
465 ) -> Result<GtResult, String> {
466 let GtAccumulator {
467 draft_path,
468 mut writer,
469 dup_masters,
470 masters_with_exact_copy,
471 n_base_masters,
472 ..
473 } = self;
474 writer
475 .finish()
476 .map_err(|e| format!("finish gt draft: {e}"))?;
477 drop(writer);
478
479 let final_schema = Arc::new(
480 Schema::new(vec![
481 Field::new("record_id", DataType::Utf8, false),
482 Field::new("master_id", DataType::Utf8, false),
483 Field::new("entity_type", low_cardinality_dict_type(), false),
484 Field::new("match_type", low_cardinality_dict_type(), false),
485 Field::new("difficulty", low_cardinality_dict_type(), false),
486 ])
487 .with_metadata(metadata.clone()),
488 );
489
490 let draft_file = std::fs::File::open(&draft_path)
491 .map_err(|e| format!("reopen gt draft {draft_path}: {e}"))?;
492 let reader = arrow::ipc::reader::FileReader::try_new(draft_file, None)
493 .map_err(|e| format!("gt draft reader: {e}"))?;
494
495 let mut sink = GtSink::new(output_format, final_path, &final_schema, metadata)?;
496
497 let mut n_exact_dup = 0usize;
498 let mut n_fuzzy_dup = 0usize;
499 let mut n_hard_neg = 0usize;
500 let mut n_unique = 0usize;
501 // Cluster membership for duplicated masters (base + duplicate copies,
502 // exact and fuzzy alike), for emitting duplicate-cluster edges after
503 // this pass. Accumulated as a flat `Vec` of packed triples (see
504 // `ClusterCsr`) instead of a `HashMap<String, Vec<...>>` -- avoids a
505 // `String` allocation per record_id/master_id, the dominant RAM cost
506 // of this bookkeeping at 200M+ records.
507 let mut cluster_pairs: Vec<(u64, u64)> = Vec::new();
508 // `difficulty` is constant for the whole run and `n` only ever takes
509 // two values across all batches (the draft's batch size, and the
510 // final remainder) — cache the built column by `n` instead of
511 // reallocating+refilling it every batch (same pattern as
512 // `pipeline.rs`'s `const_arr_cache`).
513 let mut diff_arr_cache: Option<(usize, ArrayRef)> = None;
514 // Fixed, fully-known value sets — built once, before any batch is
515 // written, and reused for every batch's `DictionaryArray` (hunt1808/
516 // H11; see `pipeline::DictValues`'s doc comment for why a *shared*
517 // dictionary is required for IPC output, not just a same-content
518 // one built fresh per batch).
519 let match_type_dict = crate::pipeline::DictValues::new([
520 "hard_neg",
521 "canary",
522 "exact_dup",
523 "fuzzy_dup",
524 "unique",
525 ]);
526 let difficulty_dict = crate::pipeline::DictValues::new([difficulty.to_string()]);
527
528 for batch_result in reader {
529 let batch = batch_result.map_err(|e| format!("read gt draft batch: {e}"))?;
530 let n = batch.num_rows();
531 let mid_col = batch.column(1).as_string::<i32>();
532 let rid_col = batch.column(0).as_string::<i32>();
533 let ident_col = batch.column(3).as_boolean();
534 let base_col = batch.column(4).as_boolean();
535
536 let mut mt_keys = arrow::array::Int32Builder::with_capacity(n);
537 for i in 0..n {
538 let mid = if mid_col.is_null(i) {
539 ""
540 } else {
541 mid_col.value(i)
542 };
543 let is_base = !base_col.is_null(i) && base_col.value(i);
544 let is_hn = mid.starts_with("HN-");
545 let is_canary = !is_hn && mid.starts_with("CANARY-");
546 // `dup_masters`/`masters_with_exact_copy` are keyed on
547 // `pack_master_key`, not the raw string (hunt3007/H3) --
548 // never `Some` for HN-/CANARY- ids, consistent with
549 // `push_dup_batch` never inserting those.
550 let master_key = if is_hn || is_canary {
551 None
552 } else {
553 crate::pipeline::pack_master_key(mid)
554 };
555 // Per-row, not per-cluster: a cluster can mix a base row, a
556 // copy the noise happened not to change, and a copy that's
557 // genuinely different. A base row has no noise of its own —
558 // it reads `exact_dup` only if the cluster actually contains
559 // a byte-identical copy (`masters_with_exact_copy`), never
560 // unconditionally (that used to mislabel every duplicated
561 // master's base row `exact_dup` even when all of its copies
562 // were fuzzy).
563 let is_identical = if is_base {
564 master_key.is_some_and(|k| masters_with_exact_copy.contains(&k))
565 } else {
566 !ident_col.is_null(i) && ident_col.value(i)
567 };
568 let is_dup_master = master_key.is_some_and(|k| dup_masters.contains(&k));
569 let mt = if is_hn {
570 n_hard_neg += 1;
571 "hard_neg"
572 } else if is_canary {
573 "canary"
574 } else if is_dup_master {
575 if is_identical {
576 n_exact_dup += 1;
577 "exact_dup"
578 } else {
579 n_fuzzy_dup += 1;
580 "fuzzy_dup"
581 }
582 } else {
583 n_unique += 1;
584 "unique"
585 };
586 mt_keys.append_value(match_type_dict.key(mt));
587 // Every row of a duplicated master belongs to its cluster,
588 // tagged with its own identical/fuzzy status. Only tracked
589 // when `--graph` is enabled: `cluster_pairs` is exclusively
590 // consumed by `graph_gen::push_dup_clusters` to emit
591 // exact_dup/fuzzy_dup edges, so building it when there's no
592 // graph output to write is pure wasted RAM -- one entry per
593 // duplicate-cluster row (base + every copy), which at
594 // 100M+ records with a low singleton fraction (hell tier)
595 // can be the majority of the dataset.
596 if track_clusters && is_dup_master {
597 let master_key = master_key.expect("is_dup_master implies master_key is Some");
598 let record_idx = crate::pipeline::parse_record_idx(rid_col.value(i))
599 .unwrap_or_else(|| {
600 panic!(
601 "cluster_map: record_id {:?} isn't the fixed \
602 \"R-\" + 13-digit shape produced by \
603 `record_id_string`",
604 rid_col.value(i)
605 )
606 });
607 cluster_pairs.push((master_key, pack_ridx_identical(record_idx, is_identical)));
608 }
609 }
610 let mt_arr = match_type_dict.finish_keys(mt_keys.finish());
611 let diff_arr = match &diff_arr_cache {
612 Some((cached_n, arr)) if *cached_n == n => arr.clone(),
613 _ => {
614 let arr = difficulty_dict.const_array(difficulty, n);
615 diff_arr_cache = Some((n, arr.clone()));
616 arr
617 }
618 };
619
620 let final_batch = RecordBatch::try_new(
621 final_schema.clone(),
622 vec![
623 batch.column(0).clone(),
624 batch.column(1).clone(),
625 batch.column(2).clone(),
626 mt_arr,
627 diff_arr,
628 ],
629 )
630 .map_err(|e| format!("build final gt batch: {e}"))?;
631
632 sink.write(&final_batch)?;
633 }
634
635 sink.finish()?;
636 std::fs::remove_file(&draft_path).ok();
637
638 Ok(GtResult {
639 n_exact_dup,
640 n_fuzzy_dup,
641 n_hard_neg,
642 n_unique,
643 n_masters: n_base_masters,
644 cluster_map: ClusterCsr::build(cluster_pairs),
645 })
646 }
647}
648
649enum GtSink {
650 Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
651 Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
652}
653
654impl GtSink {
655 fn new(
656 output_format: &str,
657 path: &str,
658 schema: &Arc<Schema>,
659 metadata: &HashMap<String, String>,
660 ) -> Result<Self, String> {
661 let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
662 if output_format == "parquet" {
663 use parquet::basic::{Compression, ZstdLevel};
664 use parquet::file::properties::WriterProperties;
665 let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
666 let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
667 .iter()
668 .map(|(k, v)| parquet::file::metadata::KeyValue {
669 key: k.clone(),
670 value: Some(v.clone()),
671 })
672 .collect();
673 let props = WriterProperties::builder()
674 .set_compression(Compression::ZSTD(zstd))
675 .set_data_page_size_limit(1_048_576)
676 .set_key_value_metadata(Some(meta_kv))
677 .build();
678 let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
679 .map_err(|e| format!("gt parquet writer: {e}"))?;
680 Ok(GtSink::Parquet(Box::new(writer)))
681 } else {
682 let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
683 .map_err(|e| format!("gt ipc writer: {e}"))?;
684 Ok(GtSink::Ipc(Box::new(writer)))
685 }
686 }
687
688 fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
689 match self {
690 GtSink::Ipc(w) => w.write(batch).map_err(|e| format!("write gt ipc: {e}")),
691 GtSink::Parquet(w) => w.write(batch).map_err(|e| format!("write gt parquet: {e}")),
692 }
693 }
694
695 fn finish(self) -> Result<(), String> {
696 match self {
697 GtSink::Ipc(mut w) => w.finish().map_err(|e| format!("finish gt ipc: {e}")),
698 GtSink::Parquet(w) => w
699 .close()
700 .map(|_| ())
701 .map_err(|e| format!("close gt parquet: {e}")),
702 }
703 }
704}
705
706#[cfg(test)]
707mod tests {
708 use super::*;
709 use arrow::array::StringArray;
710
711 fn tmp_path(name: &str) -> String {
712 let mut p = std::env::temp_dir();
713 p.push(format!(
714 "dupehell_gt_test_{name}_{}.ipc",
715 std::process::id()
716 ));
717 p.to_string_lossy().into_owned()
718 }
719
720 fn arr(values: Vec<&str>) -> ArrayRef {
721 Arc::new(StringArray::from(values))
722 }
723
724 /// Same as `arr`, but for computed (owned) strings -- used to build
725 /// `record_id`/`master_id` fixtures matching the exact fixed-width shape
726 /// `ClusterCsr`'s packing relies on (`pipeline::record_id_string`/
727 /// `entity_prefix`/`pad_string`), which a bare string literal like
728 /// `"R2"` doesn't.
729 fn arr_owned(values: Vec<String>) -> ArrayRef {
730 Arc::new(StringArray::from(values))
731 }
732
733 fn barr(values: Vec<bool>) -> ArrayRef {
734 Arc::new(BooleanArray::from(values))
735 }
736
737 /// `entity_type` fixtures (hunt1808/H11): the real column is
738 /// `Dictionary(Int32, Utf8)`, built from `add_metadata_and_align`'s
739 /// `entity_type_dict` — these tests build the draft schema directly,
740 /// so they need to hand `push_*_batch` a matching dictionary array
741 /// rather than the plain `StringArray` `arr()` builds. `dict` must be
742 /// the SAME `DictValues` instance across every `dict_arr` call feeding
743 /// one draft file — a fresh dictionary per call is exactly the
744 /// "Dictionary replacement" bug this whole hunt/H11 fixes in the real
745 /// pipeline code, and IPC enforces it just as strictly in tests.
746 fn dict_arr(dict: &crate::pipeline::DictValues, values: Vec<&str>) -> ArrayRef {
747 let keys =
748 arrow::array::Int32Array::from(values.iter().map(|v| dict.key(v)).collect::<Vec<_>>());
749 dict.finish_keys(keys)
750 }
751
752 /// Reads the final GT file into a `record_id -> match_type` map. A map
753 /// (rather than a positional `Vec`) keeps the assertions independent of
754 /// draft row order, which is caller-determined (base/dup/other batches
755 /// are pushed as separate, non-interleaved groups in the real pipeline).
756 fn read_match_types(final_path: &str) -> HashMap<String, String> {
757 let file = std::fs::File::open(final_path).unwrap();
758 let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
759 let mut out = HashMap::new();
760 for batch in reader {
761 let batch = batch.unwrap();
762 let rid = batch
763 .column_by_name("record_id")
764 .unwrap()
765 .as_string::<i32>();
766 // `match_type` is `Dictionary(Int32, Utf8)` (hunt1808/H11) —
767 // cast back to plain `Utf8` for this test-only comparison.
768 let mt_col =
769 arrow::compute::cast(batch.column_by_name("match_type").unwrap(), &DataType::Utf8)
770 .unwrap();
771 let mt = mt_col.as_string::<i32>();
772 for i in 0..batch.num_rows() {
773 out.insert(rid.value(i).to_string(), mt.value(i).to_string());
774 }
775 }
776 out
777 }
778
779 /// Builds the fixed-width `"{entity_prefix}-{pad_string}"` master_id
780 /// shape `ClusterCsr`'s packing relies on -- all these fixture-building
781 /// tests use a single entity index (0).
782 fn mid(n: usize) -> String {
783 format!(
784 "{}-{}",
785 crate::pipeline::entity_prefix(0),
786 crate::pipeline::pad_string(n)
787 )
788 }
789
790 /// Alias for `pipeline::record_id_string`, for brevity in fixtures.
791 fn rid(i: usize) -> String {
792 crate::pipeline::record_id_string(i)
793 }
794
795 #[test]
796 fn test_gt_accumulator_basic() {
797 let draft = tmp_path("basic_draft");
798 let final_path = tmp_path("basic_final");
799
800 let mut acc = GtAccumulator::new(&draft).unwrap();
801 let et_dict = crate::pipeline::DictValues::new(["person"]);
802 // Base rows: rid(1) (master mid(1), later duplicated), rid(3)
803 // (singleton mid(2)), rid(6) (singleton mid(5)).
804 acc.push_base_batch(
805 &arr_owned(vec![rid(1), rid(3), rid(6)]),
806 &dict_arr(&et_dict, vec!["person", "person", "person"]),
807 &arr_owned(vec![mid(1), mid(2), mid(5)]),
808 )
809 .unwrap();
810 // Dup row: rid(2) duplicates mid(1), unchanged by its noise pass.
811 acc.push_dup_batch(
812 &arr_owned(vec![rid(2)]),
813 &dict_arr(&et_dict, vec!["person"]),
814 &arr_owned(vec![mid(1)]),
815 &barr(vec![true]),
816 )
817 .unwrap();
818 // Hard negatives: rid(4), rid(5).
819 acc.push_other_batch(
820 &arr_owned(vec![rid(4), rid(5)]),
821 &dict_arr(&et_dict, vec!["person", "person"]),
822 &arr(vec!["HN-0000003", "HN-0000004"]),
823 )
824 .unwrap();
825
826 let GtResult {
827 n_exact_dup: ed,
828 n_hard_neg: hn,
829 n_unique: un,
830 n_masters: masters,
831 ..
832 } = acc
833 .finish("medium", "ipc", &final_path, &HashMap::new(), false)
834 .unwrap();
835 assert_eq!(ed, 2);
836 assert_eq!(hn, 2);
837 assert_eq!(un, 2);
838 assert_eq!(masters, 3); // mid(1), mid(2), mid(5)
839
840 let match_types = read_match_types(&final_path);
841 assert_eq!(match_types[&rid(1)], "exact_dup");
842 assert_eq!(match_types[&rid(2)], "exact_dup");
843 assert_eq!(match_types[&rid(3)], "unique");
844 assert_eq!(match_types[&rid(4)], "hard_neg");
845 assert_eq!(match_types[&rid(5)], "hard_neg");
846 assert_eq!(match_types[&rid(6)], "unique");
847
848 std::fs::remove_file(&final_path).ok();
849 }
850
851 /// A duplicate copy whose noise pass actually changed something
852 /// (`is_identical = false`) must be classified `fuzzy_dup`, not
853 /// `exact_dup` — the master itself, and any sibling copy that was left
854 /// unchanged, keep `exact_dup` independently, per row.
855 #[test]
856 fn test_gt_accumulator_fuzzy_dup() {
857 let draft = tmp_path("fuzzy_draft");
858 let final_path = tmp_path("fuzzy_final");
859
860 let mut acc = GtAccumulator::new(&draft).unwrap();
861 let et_dict = crate::pipeline::DictValues::new(["person"]);
862 // Triplet: rid(1) is the master, rid(2) was left unchanged by its
863 // noise pass, rid(3) was genuinely altered.
864 acc.push_base_batch(
865 &arr_owned(vec![rid(1)]),
866 &dict_arr(&et_dict, vec!["person"]),
867 &arr_owned(vec![mid(1)]),
868 )
869 .unwrap();
870 acc.push_dup_batch(
871 &arr_owned(vec![rid(2), rid(3)]),
872 &dict_arr(&et_dict, vec!["person", "person"]),
873 &arr_owned(vec![mid(1), mid(1)]),
874 &barr(vec![true, false]),
875 )
876 .unwrap();
877
878 let GtResult {
879 n_exact_dup: ed,
880 n_fuzzy_dup: fd,
881 n_unique: un,
882 n_masters: masters,
883 ..
884 } = acc
885 .finish("medium", "ipc", &final_path, &HashMap::new(), false)
886 .unwrap();
887 assert_eq!(ed, 2); // rid(1) (master) + rid(2) (unchanged copy)
888 assert_eq!(fd, 1); // rid(3) (genuinely noised copy)
889 assert_eq!(un, 0);
890 assert_eq!(masters, 1);
891
892 let match_types = read_match_types(&final_path);
893 assert_eq!(match_types[&rid(1)], "exact_dup");
894 assert_eq!(match_types[&rid(2)], "exact_dup");
895 assert_eq!(match_types[&rid(3)], "fuzzy_dup");
896
897 std::fs::remove_file(&final_path).ok();
898 }
899
900 #[test]
901 fn test_gt_accumulator_multi_batch() {
902 let draft = tmp_path("multi_draft");
903 let final_path = tmp_path("multi_final");
904
905 let mut acc = GtAccumulator::new(&draft).unwrap();
906 let et_dict = crate::pipeline::DictValues::new(["person"]);
907 // Base batch: two masters, one of which (mid(1)) is duplicated in a
908 // later, separate dup batch — simulating a master's duplicate
909 // landing in a different batch than its base row.
910 acc.push_base_batch(
911 &arr_owned(vec![rid(1), rid(3)]),
912 &dict_arr(&et_dict, vec!["person", "person"]),
913 &arr_owned(vec![mid(1), mid(2)]),
914 )
915 .unwrap();
916 acc.push_dup_batch(
917 &arr_owned(vec![rid(2)]),
918 &dict_arr(&et_dict, vec!["person"]),
919 &arr_owned(vec![mid(1)]),
920 &barr(vec![true]),
921 )
922 .unwrap();
923
924 let GtResult {
925 n_exact_dup: ed,
926 n_hard_neg: hn,
927 n_unique: un,
928 n_masters: masters,
929 ..
930 } = acc
931 .finish("medium", "ipc", &final_path, &HashMap::new(), false)
932 .unwrap();
933 assert_eq!(ed, 2); // both rows of M-0000001
934 assert_eq!(hn, 0);
935 assert_eq!(un, 1);
936 assert_eq!(masters, 2);
937
938 std::fs::remove_file(&final_path).ok();
939 }
940
941 /// Regression: two different entity types can produce master_ids that
942 /// share the same numeric suffix (each entity's index restarts at 0),
943 /// but must not be counted as duplicates of each other.
944 #[test]
945 fn test_no_cross_entity_suffix_collision() {
946 let draft = tmp_path("suffix_draft");
947 let final_path = tmp_path("suffix_final");
948
949 let mut acc = GtAccumulator::new(&draft).unwrap();
950 let et_dict = crate::pipeline::DictValues::new(["person", "account"]);
951 acc.push_base_batch(
952 &arr(vec!["R1", "R2"]),
953 &dict_arr(&et_dict, vec!["person", "account"]),
954 &arr(vec!["PERSON-0000001", "ACCOUNT-0000001"]),
955 )
956 .unwrap();
957
958 let GtResult {
959 n_exact_dup: ed,
960 n_unique: un,
961 n_masters: masters,
962 ..
963 } = acc
964 .finish("medium", "ipc", &final_path, &HashMap::new(), false)
965 .unwrap();
966 assert_eq!(ed, 0);
967 assert_eq!(un, 2);
968 assert_eq!(masters, 2);
969
970 std::fs::remove_file(&final_path).ok();
971 }
972
973 /// `finish()` returns a `cluster_map` mapping each duplicated master_id
974 /// to the full set of `(record_id, is_identical)` pairs in its cluster
975 /// (base + duplicate copies). This is the structure consumed by
976 /// `push_dup_clusters` to decide, per edge, `exact_dup` vs `fuzzy_dup`.
977 #[test]
978 fn test_cluster_map_contents() {
979 let draft = tmp_path("cm_draft");
980 let final_path = tmp_path("cm_final");
981
982 let mut acc = GtAccumulator::new(&draft).unwrap();
983 let et_dict = crate::pipeline::DictValues::new(["person"]);
984 acc.push_base_batch(
985 &arr_owned(vec![rid(1), rid(3), rid(6)]),
986 &dict_arr(&et_dict, vec!["person", "person", "person"]),
987 &arr_owned(vec![mid(1), mid(2), mid(5)]),
988 )
989 .unwrap();
990 // rid(2) (mid(1)) stayed identical; rid(7) (mid(5)) was genuinely noised.
991 acc.push_dup_batch(
992 &arr_owned(vec![rid(2), rid(7)]),
993 &dict_arr(&et_dict, vec!["person", "person"]),
994 &arr_owned(vec![mid(1), mid(5)]),
995 &barr(vec![true, false]),
996 )
997 .unwrap();
998 acc.push_other_batch(
999 &arr_owned(vec![rid(4), rid(5)]),
1000 &dict_arr(&et_dict, vec!["person", "person"]),
1001 &arr(vec!["HN-0000003", "HN-0000004"]),
1002 )
1003 .unwrap();
1004
1005 let GtResult {
1006 cluster_map: cm, ..
1007 } = acc
1008 .finish("medium", "ipc", &final_path, &HashMap::new(), true)
1009 .unwrap();
1010
1011 // Only duplicated masters appear; singletons (mid(2)) do not. Groups
1012 // come out in ascending packed-master-key order, i.e. mid(1) (local
1013 // index 1) before mid(5) (local index 5).
1014 assert_eq!(cm.n_clusters(), 2);
1015 let mut groups = cm.groups();
1016 let (records1, idents1) = groups.next().unwrap();
1017 assert_eq!(records1, &[1, 2]);
1018 assert_eq!(idents1.to_vec(), vec![true, true]);
1019 // rid(6) is the base of mid(5); its only copy (rid(7)) was genuinely
1020 // noised, so the cluster has no byte-identical pair at all — rid(6)
1021 // must NOT read as identical just because it's the base (that was
1022 // the bug: base rows used to hardcode `is_identical = true`).
1023 let (records5, idents5) = groups.next().unwrap();
1024 assert_eq!(records5, &[6, 7]);
1025 assert_eq!(idents5.to_vec(), vec![false, false]);
1026 assert!(groups.next().is_none());
1027
1028 std::fs::remove_file(&final_path).ok();
1029 }
1030
1031 // ── perf-hunt RAM pass (hunt1708.md) ─────────────────────────────────
1032 //
1033 // Not a correctness test -- #[ignore]'d so it never runs in normal
1034 // `cargo test`. Isolated measurement (never a full pipeline run, per
1035 // perf-hunt rule 4): compares process RSS growth for two separate
1036 // `FxHashSet<u64>` (`dup_masters` + `masters_with_exact_copy` as they
1037 // exist today) against one combined `FxHashMap<u64, bool>`, at the
1038 // cardinality actually observed on a real run (aviation/hell/50M:
1039 // masters=21,666,666, uniques=5,855,236 -> dup_masters.len() ~=
1040 // 15,811,430; masters_with_exact_copy is a subset, sized here at 70%
1041 // of that as a representative estimate -- exact_dups=23,237,057 of
1042 // fuzzy+exact=44,144,764 total duplicate rows is ~53% at the row
1043 // level, but a master with ANY identical copy is more common than
1044 // that row-level fraction suggests, so 70% is a deliberately
1045 // generous (not cherry-picked-low) estimate for this comparison).
1046 // Run with:
1047 // cargo test --release gt::tests::hunt_ram_dup_masters_vs_combined_map -- --ignored --nocapture
1048 #[test]
1049 #[ignore]
1050 fn hunt_ram_dup_masters_vs_combined_map() {
1051 fn rss_mb() -> f64 {
1052 let pid = sysinfo::Pid::from_u32(std::process::id());
1053 let mut sys = sysinfo::System::new();
1054 sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
1055 sys.process(pid)
1056 .map(|p| p.memory() as f64 / 1e6)
1057 .unwrap_or(0.0)
1058 }
1059
1060 const N_DUP_MASTERS: usize = 15_811_430;
1061 const N_EXACT_COPY: usize = (N_DUP_MASTERS as f64 * 0.70) as usize;
1062
1063 // Deterministic, well-spread u64 keys (not sequential -- sequential
1064 // keys would let the hash table's bucket layout be unrealistically
1065 // cache-friendly compared to the real `pack_master_key` output,
1066 // which mixes an entity_idx high-bit prefix with a local counter).
1067 // A cheap SplitMix64-style mix is enough to spread bits without
1068 // pulling in a new dependency for a throwaway measurement.
1069 fn mix(i: u64) -> u64 {
1070 let mut z = i.wrapping_add(0x9E3779B97F4A7C15);
1071 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1072 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1073 z ^ (z >> 31)
1074 }
1075
1076 // ── Scenario A: two separate FxHashSet<u64> (today's shape) ────────
1077 let rss0 = rss_mb();
1078 let mut dup_masters: rustc_hash::FxHashSet<u64> =
1079 rustc_hash::FxHashSet::with_capacity_and_hasher(N_DUP_MASTERS, Default::default());
1080 for i in 0..N_DUP_MASTERS as u64 {
1081 dup_masters.insert(mix(i));
1082 }
1083 let mut masters_with_exact_copy: rustc_hash::FxHashSet<u64> =
1084 rustc_hash::FxHashSet::with_capacity_and_hasher(N_EXACT_COPY, Default::default());
1085 for i in 0..N_EXACT_COPY as u64 {
1086 // Subset of the same key space as dup_masters (every master
1087 // with an exact copy is also in dup_masters), same mix.
1088 masters_with_exact_copy.insert(mix(i));
1089 }
1090 let rss_two_sets = rss_mb();
1091 println!(
1092 "[hunt_ram] two FxHashSet<u64> ({} + {} entries): rss delta = {:.1} Mo (rss0={:.1} -> {:.1})",
1093 dup_masters.len(),
1094 masters_with_exact_copy.len(),
1095 rss_two_sets - rss0,
1096 rss0,
1097 rss_two_sets
1098 );
1099 drop(dup_masters);
1100 drop(masters_with_exact_copy);
1101
1102 // ── Scenario B: one combined FxHashMap<u64, bool> ──────────────────
1103 let rss1 = rss_mb();
1104 let mut combined: rustc_hash::FxHashMap<u64, bool> =
1105 rustc_hash::FxHashMap::with_capacity_and_hasher(N_DUP_MASTERS, Default::default());
1106 for i in 0..N_DUP_MASTERS as u64 {
1107 combined.insert(mix(i), i < N_EXACT_COPY as u64);
1108 }
1109 let rss_combined = rss_mb();
1110 println!(
1111 "[hunt_ram] one FxHashMap<u64,bool> ({} entries): rss delta = {:.1} Mo (rss1={:.1} -> {:.1})",
1112 combined.len(),
1113 rss_combined - rss1,
1114 rss1,
1115 rss_combined
1116 );
1117 drop(combined);
1118
1119 let delta_two = rss_two_sets - rss0;
1120 let delta_one = rss_combined - rss1;
1121 println!(
1122 "[hunt_ram] combined map saves {:.1} Mo ({:.1}%) vs two separate sets, at this cardinality",
1123 delta_two - delta_one,
1124 100.0 * (delta_two - delta_one) / delta_two.max(1.0)
1125 );
1126 }
1127}