1use std::collections::HashMap;
2
3use zer_core::{
4 record::{Record, RecordId},
5 schema::Schema,
6 traits::{BlockIndex, Blocker},
7};
8
9use crate::keys::BlockingKey;
10
11pub struct CompositeBlocker {
19 keys: Vec<Box<dyn BlockingKey>>,
20 source_remaps: HashMap<String, HashMap<String, String>>,
21}
22
23impl CompositeBlocker {
24 pub fn new() -> Self {
25 Self {
26 keys: vec![],
27 source_remaps: HashMap::new(),
28 }
29 }
30
31 #[allow(clippy::should_implement_trait)]
32 pub fn add(mut self, key: impl BlockingKey + 'static) -> Self {
33 self.keys.push(Box::new(key));
34 self
35 }
36
37 pub fn add_boxed(mut self, key: Box<dyn BlockingKey>) -> Self {
38 self.keys.push(key);
39 self
40 }
41
42 pub fn with_source_remap(
47 mut self,
48 source: impl Into<String>,
49 remap: HashMap<String, String>,
50 ) -> Self {
51 self.source_remaps.insert(source.into(), remap);
52 self
53 }
54
55 fn effective_record(&self, record: &Record) -> Option<Record> {
56 let src = record.source.as_deref()?;
57 let remap = self.source_remaps.get(src)?;
58 let mut new_rec = Record::new(record.id);
59 if let Some(s) = &record.source {
60 new_rec = new_rec.with_source(s);
61 }
62 for (field_name, value) in &record.fields {
63 let canonical = remap
64 .get(field_name)
65 .cloned()
66 .unwrap_or_else(|| field_name.clone());
67 new_rec.fields.insert(canonical, value.clone());
68 }
69 Some(new_rec)
70 }
71}
72
73impl Default for CompositeBlocker {
74 fn default() -> Self {
75 Self::new()
76 }
77}
78
79impl Blocker for CompositeBlocker {
80 fn blocking_keys(&self, record: &Record, schema: &Schema) -> Vec<String> {
81 let remapped = self.effective_record(record);
82 let effective = remapped.as_ref().unwrap_or(record);
83 self.keys
84 .iter()
85 .flat_map(|k| {
86 k.extract(effective, schema)
87 .into_iter()
88 .map(|val| format!("{}:{}", k.name(), val))
89 })
90 .collect()
91 }
92
93 fn index_record(&self, record: &Record, schema: &Schema, index: &mut dyn BlockIndex) {
94 let keys = self.blocking_keys(record, schema);
95 index.insert(record.id, keys);
96 }
97
98 fn candidates(
99 &self,
100 record: &Record,
101 schema: &Schema,
102 index: &dyn BlockIndex,
103 ) -> Vec<RecordId> {
104 let keys = self.blocking_keys(record, schema);
105 index.lookup_union(&keys, record.id)
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112 use crate::{index::InvertedIndex, keys::ExactFieldKey};
113 use zer_core::{
114 record::FieldValue,
115 schema::{FieldKind, SchemaBuilder},
116 };
117
118 fn schema() -> Schema {
119 SchemaBuilder::new()
120 .field("category", FieldKind::Categorical)
121 .build()
122 .unwrap()
123 }
124
125 #[test]
126 fn index_and_candidates_round_trip() {
127 let schema = schema();
128 let blocker = CompositeBlocker::new().add(ExactFieldKey::new("category"));
129 let mut idx = InvertedIndex::new();
130
131 let r1 = Record::new(1).insert("category", FieldValue::Text("TypeA".into()));
132 let r2 = Record::new(2).insert("category", FieldValue::Text("TypeA".into()));
133 let r3 = Record::new(3).insert("category", FieldValue::Text("TypeB".into()));
134
135 blocker.index_record(&r1, &schema, &mut idx);
136 blocker.index_record(&r2, &schema, &mut idx);
137 blocker.index_record(&r3, &schema, &mut idx);
138
139 let cands_r1 = blocker.candidates(&r1, &schema, &idx);
140 assert!(cands_r1.contains(&2), "r2 should be a candidate for r1");
141 assert!(!cands_r1.contains(&1), "r1 should not be its own candidate");
142 assert!(
143 !cands_r1.contains(&3),
144 "r3 should not match r1 (different category)"
145 );
146 }
147
148 #[test]
149 fn no_self_candidates() {
150 let schema = schema();
151 let blocker = CompositeBlocker::new().add(ExactFieldKey::new("category"));
152 let mut idx = InvertedIndex::new();
153
154 let r = Record::new(1).insert("category", FieldValue::Text("X".into()));
155 blocker.index_record(&r, &schema, &mut idx);
156
157 let cands = blocker.candidates(&r, &schema, &idx);
158 assert!(!cands.contains(&1));
159 }
160}