Skip to main content

zer_blocking/
blocker.rs

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
11/// Composite blocker that applies multiple blocking keys.
12///
13/// For cross-schema linkage, `source_remaps` stores per-source field-name
14/// translations (b_field to a_field).  Before key extraction, any record
15/// whose `source` label has an entry in `source_remaps` gets its fields
16/// renamed to the canonical (A-side) names so the existing `BlockingKey`
17/// implementations can extract values without knowing about schema differences.
18pub 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    /// Register a field-name remap for records from `source`.
43    ///
44    /// `remap` maps b_field to a_field so that the source-B fields are
45    /// visible under canonical source-A names during blocking key extraction.
46    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}