use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, StringArray, StringBuilder};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::collections::HashMap;
use std::sync::Arc;
pub struct GtResult {
pub n_exact_dup: usize,
pub n_fuzzy_dup: usize,
pub n_hard_neg: usize,
pub n_unique: usize,
pub n_masters: usize,
pub cluster_map: HashMap<String, Vec<(String, bool)>>,
}
fn draft_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("record_id", DataType::Utf8, false),
Field::new("master_id", DataType::Utf8, false),
Field::new("entity_type", DataType::Utf8, false),
Field::new("is_identical", DataType::Boolean, false),
Field::new("is_base", DataType::Boolean, false),
]))
}
pub struct GtAccumulator {
draft_path: String,
writer: arrow::ipc::writer::FileWriter<std::fs::File>,
schema: Arc<Schema>,
dup_masters: rustc_hash::FxHashSet<String>,
masters_with_exact_copy: rustc_hash::FxHashSet<String>,
n_base_masters: usize,
}
impl GtAccumulator {
pub fn new(draft_path: &str) -> Result<Self, String> {
let schema = draft_schema();
let file = std::fs::File::create(draft_path)
.map_err(|e| format!("create gt draft {draft_path}: {e}"))?;
let writer = arrow::ipc::writer::FileWriter::try_new(file, &schema)
.map_err(|e| format!("gt draft writer: {e}"))?;
Ok(Self {
draft_path: draft_path.to_string(),
writer,
schema,
dup_masters: rustc_hash::FxHashSet::default(),
masters_with_exact_copy: rustc_hash::FxHashSet::default(),
n_base_masters: 0,
})
}
fn write_draft(
&mut self,
record_ids: &ArrayRef,
entity_types: &ArrayRef,
master_ids: &ArrayRef,
is_identical: &ArrayRef,
is_base: &ArrayRef,
) -> Result<(), String> {
let batch = RecordBatch::try_new(
self.schema.clone(),
vec![
record_ids.clone(),
master_ids.clone(),
entity_types.clone(),
is_identical.clone(),
is_base.clone(),
],
)
.map_err(|e| format!("build gt draft batch: {e}"))?;
self.writer
.write(&batch)
.map_err(|e| format!("write gt draft: {e}"))
}
pub fn push_base_batch(
&mut self,
record_ids: &ArrayRef,
entity_types: &ArrayRef,
master_ids: &ArrayRef,
) -> Result<(), String> {
let mids = master_ids.as_string::<i32>();
for i in 0..mids.len() {
if !mids.is_null(i) {
self.n_base_masters += 1;
}
}
let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; mids.len()]));
let all_true: ArrayRef = Arc::new(BooleanArray::from(vec![true; mids.len()]));
self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_true)
}
pub fn push_dup_batch(
&mut self,
record_ids: &ArrayRef,
entity_types: &ArrayRef,
master_ids: &ArrayRef,
is_identical: &ArrayRef,
) -> Result<(), String> {
let mids = master_ids.as_string::<i32>();
let idents = is_identical.as_boolean();
for i in 0..mids.len() {
if !mids.is_null(i) {
let mid = mids.value(i);
if !self.dup_masters.contains(mid) {
self.dup_masters.insert(mid.to_string());
}
if !idents.is_null(i)
&& idents.value(i)
&& !self.masters_with_exact_copy.contains(mid)
{
self.masters_with_exact_copy.insert(mid.to_string());
}
}
}
let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; mids.len()]));
self.write_draft(
record_ids,
entity_types,
master_ids,
is_identical,
&all_false,
)
}
pub fn push_other_batch(
&mut self,
record_ids: &ArrayRef,
entity_types: &ArrayRef,
master_ids: &ArrayRef,
) -> Result<(), String> {
let all_false: ArrayRef = Arc::new(BooleanArray::from(vec![false; record_ids.len()]));
self.write_draft(record_ids, entity_types, master_ids, &all_false, &all_false)
}
pub fn finish(
self,
difficulty: &str,
output_format: &str,
final_path: &str,
metadata: &HashMap<String, String>,
) -> Result<GtResult, String> {
let GtAccumulator {
draft_path,
mut writer,
dup_masters,
masters_with_exact_copy,
n_base_masters,
..
} = self;
writer
.finish()
.map_err(|e| format!("finish gt draft: {e}"))?;
drop(writer);
let final_schema = Arc::new(
Schema::new(vec![
Field::new("record_id", DataType::Utf8, false),
Field::new("master_id", DataType::Utf8, false),
Field::new("entity_type", DataType::Utf8, false),
Field::new("match_type", DataType::Utf8, false),
Field::new("difficulty", DataType::Utf8, false),
])
.with_metadata(metadata.clone()),
);
let draft_file = std::fs::File::open(&draft_path)
.map_err(|e| format!("reopen gt draft {draft_path}: {e}"))?;
let reader = arrow::ipc::reader::FileReader::try_new(draft_file, None)
.map_err(|e| format!("gt draft reader: {e}"))?;
let mut sink = GtSink::new(output_format, final_path, &final_schema, metadata)?;
let mut n_exact_dup = 0usize;
let mut n_fuzzy_dup = 0usize;
let mut n_hard_neg = 0usize;
let mut n_unique = 0usize;
let mut cluster_map: HashMap<String, Vec<(String, bool)>> = HashMap::new();
for batch_result in reader {
let batch = batch_result.map_err(|e| format!("read gt draft batch: {e}"))?;
let n = batch.num_rows();
let mid_col = batch.column(1).as_string::<i32>();
let rid_col = batch.column(0).as_string::<i32>();
let ident_col = batch.column(3).as_boolean();
let base_col = batch.column(4).as_boolean();
let mut mt_builder = StringBuilder::with_capacity(n, n * 10);
for i in 0..n {
let mid = if mid_col.is_null(i) {
""
} else {
mid_col.value(i)
};
let is_base = !base_col.is_null(i) && base_col.value(i);
let is_identical = if is_base {
masters_with_exact_copy.contains(mid)
} else {
!ident_col.is_null(i) && ident_col.value(i)
};
let mt = if mid.starts_with("HN-") {
n_hard_neg += 1;
"hard_neg"
} else if mid.starts_with("CANARY-") {
"canary"
} else if dup_masters.contains(mid) {
if is_identical {
n_exact_dup += 1;
"exact_dup"
} else {
n_fuzzy_dup += 1;
"fuzzy_dup"
}
} else {
n_unique += 1;
"unique"
};
mt_builder.append_value(mt);
if dup_masters.contains(mid) {
cluster_map
.entry(mid.to_string())
.or_default()
.push((rid_col.value(i).to_string(), is_identical));
}
}
let mt_arr: ArrayRef = Arc::new(mt_builder.finish());
let diff_arr: ArrayRef = Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
difficulty, n,
)));
let final_batch = RecordBatch::try_new(
final_schema.clone(),
vec![
batch.column(0).clone(),
batch.column(1).clone(),
batch.column(2).clone(),
mt_arr,
diff_arr,
],
)
.map_err(|e| format!("build final gt batch: {e}"))?;
sink.write(&final_batch)?;
}
sink.finish()?;
std::fs::remove_file(&draft_path).ok();
Ok(GtResult {
n_exact_dup,
n_fuzzy_dup,
n_hard_neg,
n_unique,
n_masters: n_base_masters,
cluster_map,
})
}
}
enum GtSink {
Ipc(Box<arrow::ipc::writer::FileWriter<std::fs::File>>),
Parquet(Box<parquet::arrow::ArrowWriter<std::fs::File>>),
}
impl GtSink {
fn new(
output_format: &str,
path: &str,
schema: &Arc<Schema>,
metadata: &HashMap<String, String>,
) -> Result<Self, String> {
let file = std::fs::File::create(path).map_err(|e| format!("create {path}: {e}"))?;
if output_format == "parquet" {
use parquet::basic::{Compression, ZstdLevel};
use parquet::file::properties::WriterProperties;
let zstd = ZstdLevel::try_new(3).map_err(|e| format!("zstd: {e}"))?;
let meta_kv: Vec<parquet::file::metadata::KeyValue> = metadata
.iter()
.map(|(k, v)| parquet::file::metadata::KeyValue {
key: k.clone(),
value: Some(v.clone()),
})
.collect();
let props = WriterProperties::builder()
.set_compression(Compression::ZSTD(zstd))
.set_data_page_size_limit(1_048_576)
.set_key_value_metadata(Some(meta_kv))
.build();
let writer = parquet::arrow::ArrowWriter::try_new(file, schema.clone(), Some(props))
.map_err(|e| format!("gt parquet writer: {e}"))?;
Ok(GtSink::Parquet(Box::new(writer)))
} else {
let writer = arrow::ipc::writer::FileWriter::try_new(file, schema)
.map_err(|e| format!("gt ipc writer: {e}"))?;
Ok(GtSink::Ipc(Box::new(writer)))
}
}
fn write(&mut self, batch: &RecordBatch) -> Result<(), String> {
match self {
GtSink::Ipc(w) => w.write(batch).map_err(|e| format!("write gt ipc: {e}")),
GtSink::Parquet(w) => w.write(batch).map_err(|e| format!("write gt parquet: {e}")),
}
}
fn finish(self) -> Result<(), String> {
match self {
GtSink::Ipc(mut w) => w.finish().map_err(|e| format!("finish gt ipc: {e}")),
GtSink::Parquet(w) => w
.close()
.map(|_| ())
.map_err(|e| format!("close gt parquet: {e}")),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::StringArray;
fn tmp_path(name: &str) -> String {
let mut p = std::env::temp_dir();
p.push(format!(
"dupehell_gt_test_{name}_{}.ipc",
std::process::id()
));
p.to_string_lossy().into_owned()
}
fn arr(values: Vec<&str>) -> ArrayRef {
Arc::new(StringArray::from(values))
}
fn barr(values: Vec<bool>) -> ArrayRef {
Arc::new(BooleanArray::from(values))
}
fn read_match_types(final_path: &str) -> HashMap<String, String> {
let file = std::fs::File::open(final_path).unwrap();
let reader = arrow::ipc::reader::FileReader::try_new(file, None).unwrap();
let mut out = HashMap::new();
for batch in reader {
let batch = batch.unwrap();
let rid = batch
.column_by_name("record_id")
.unwrap()
.as_string::<i32>();
let mt = batch
.column_by_name("match_type")
.unwrap()
.as_string::<i32>();
for i in 0..batch.num_rows() {
out.insert(rid.value(i).to_string(), mt.value(i).to_string());
}
}
out
}
#[test]
fn test_gt_accumulator_basic() {
let draft = tmp_path("basic_draft");
let final_path = tmp_path("basic_final");
let mut acc = GtAccumulator::new(&draft).unwrap();
acc.push_base_batch(
&arr(vec!["R1", "R3", "R6"]),
&arr(vec!["person", "person", "person"]),
&arr(vec!["M-0000001", "M-0000002", "M-0000005"]),
)
.unwrap();
acc.push_dup_batch(
&arr(vec!["R2"]),
&arr(vec!["person"]),
&arr(vec!["M-0000001"]),
&barr(vec![true]),
)
.unwrap();
acc.push_other_batch(
&arr(vec!["R4", "R5"]),
&arr(vec!["person", "person"]),
&arr(vec!["HN-0000003", "HN-0000004"]),
)
.unwrap();
let GtResult {
n_exact_dup: ed,
n_hard_neg: hn,
n_unique: un,
n_masters: masters,
..
} = acc
.finish("medium", "ipc", &final_path, &HashMap::new())
.unwrap();
assert_eq!(ed, 2);
assert_eq!(hn, 2);
assert_eq!(un, 2);
assert_eq!(masters, 3);
let match_types = read_match_types(&final_path);
assert_eq!(match_types["R1"], "exact_dup");
assert_eq!(match_types["R2"], "exact_dup");
assert_eq!(match_types["R3"], "unique");
assert_eq!(match_types["R4"], "hard_neg");
assert_eq!(match_types["R5"], "hard_neg");
assert_eq!(match_types["R6"], "unique");
std::fs::remove_file(&final_path).ok();
}
#[test]
fn test_gt_accumulator_fuzzy_dup() {
let draft = tmp_path("fuzzy_draft");
let final_path = tmp_path("fuzzy_final");
let mut acc = GtAccumulator::new(&draft).unwrap();
acc.push_base_batch(
&arr(vec!["R1"]),
&arr(vec!["person"]),
&arr(vec!["M-0000001"]),
)
.unwrap();
acc.push_dup_batch(
&arr(vec!["R2", "R3"]),
&arr(vec!["person", "person"]),
&arr(vec!["M-0000001", "M-0000001"]),
&barr(vec![true, false]),
)
.unwrap();
let GtResult {
n_exact_dup: ed,
n_fuzzy_dup: fd,
n_unique: un,
n_masters: masters,
..
} = acc
.finish("medium", "ipc", &final_path, &HashMap::new())
.unwrap();
assert_eq!(ed, 2); assert_eq!(fd, 1); assert_eq!(un, 0);
assert_eq!(masters, 1);
let match_types = read_match_types(&final_path);
assert_eq!(match_types["R1"], "exact_dup");
assert_eq!(match_types["R2"], "exact_dup");
assert_eq!(match_types["R3"], "fuzzy_dup");
std::fs::remove_file(&final_path).ok();
}
#[test]
fn test_gt_accumulator_multi_batch() {
let draft = tmp_path("multi_draft");
let final_path = tmp_path("multi_final");
let mut acc = GtAccumulator::new(&draft).unwrap();
acc.push_base_batch(
&arr(vec!["R1", "R3"]),
&arr(vec!["person", "person"]),
&arr(vec!["M-0000001", "M-0000002"]),
)
.unwrap();
acc.push_dup_batch(
&arr(vec!["R2"]),
&arr(vec!["person"]),
&arr(vec!["M-0000001"]),
&barr(vec![true]),
)
.unwrap();
let GtResult {
n_exact_dup: ed,
n_hard_neg: hn,
n_unique: un,
n_masters: masters,
..
} = acc
.finish("medium", "ipc", &final_path, &HashMap::new())
.unwrap();
assert_eq!(ed, 2); assert_eq!(hn, 0);
assert_eq!(un, 1);
assert_eq!(masters, 2);
std::fs::remove_file(&final_path).ok();
}
#[test]
fn test_no_cross_entity_suffix_collision() {
let draft = tmp_path("suffix_draft");
let final_path = tmp_path("suffix_final");
let mut acc = GtAccumulator::new(&draft).unwrap();
acc.push_base_batch(
&arr(vec!["R1", "R2"]),
&arr(vec!["person", "account"]),
&arr(vec!["PERSON-0000001", "ACCOUNT-0000001"]),
)
.unwrap();
let GtResult {
n_exact_dup: ed,
n_unique: un,
n_masters: masters,
..
} = acc
.finish("medium", "ipc", &final_path, &HashMap::new())
.unwrap();
assert_eq!(ed, 0);
assert_eq!(un, 2);
assert_eq!(masters, 2);
std::fs::remove_file(&final_path).ok();
}
#[test]
fn test_cluster_map_contents() {
let draft = tmp_path("cm_draft");
let final_path = tmp_path("cm_final");
let mut acc = GtAccumulator::new(&draft).unwrap();
acc.push_base_batch(
&arr(vec!["R1", "R3", "R6"]),
&arr(vec!["person", "person", "person"]),
&arr(vec!["M-0000001", "M-0000002", "M-0000005"]),
)
.unwrap();
acc.push_dup_batch(
&arr(vec!["R2", "R7"]),
&arr(vec!["person", "person"]),
&arr(vec!["M-0000001", "M-0000005"]),
&barr(vec![true, false]),
)
.unwrap();
acc.push_other_batch(
&arr(vec!["R4", "R5"]),
&arr(vec!["person", "person"]),
&arr(vec!["HN-0000003", "HN-0000004"]),
)
.unwrap();
let GtResult {
cluster_map: cm, ..
} = acc
.finish("medium", "ipc", &final_path, &HashMap::new())
.unwrap();
assert_eq!(cm.len(), 2);
let mut m1 = cm.get("M-0000001").unwrap().clone();
m1.sort();
assert_eq!(m1, vec![("R1".to_string(), true), ("R2".to_string(), true)]);
let mut m5 = cm.get("M-0000005").unwrap().clone();
m5.sort();
assert_eq!(
m5,
vec![("R6".to_string(), false), ("R7".to_string(), false)]
);
assert!(!cm.contains_key("M-0000002"));
std::fs::remove_file(&final_path).ok();
}
}