zer_core/
traits.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3use std::sync::RwLock;
4
5use crate::{
6    comparison::{ComparisonBatch, ComparisonVector},
7    entity::{Entity, EntityId},
8    error::ZerError,
9    record::{Record, RecordId},
10    record_pool::RecordPool,
11    schema::Schema,
12    scoring::{ModelParams, ScoredPair},
13};
14
15pub type Result<T> = std::result::Result<T, ZerError>;
16
17// ── RecordStore ───────────────────────────────────────────────────────────────
18
19/// Backing store for records used during ingestion and batch runs.
20pub trait RecordStore: Send + Sync {
21    /// Persist a record.  Must be callable from the ingester background task.
22    fn insert(&self, record: Record);
23
24    /// Retrieve a single record by ID.  Returns `None` if not present.
25    fn get(&self, id: RecordId) -> Option<Cow<'_, Record>>;
26
27    /// Retrieve multiple records in one call (allows batch I/O optimisation).
28    /// The default impl calls `get` in a loop; override for bulk reads.
29    fn get_many(&self, ids: &[RecordId]) -> Vec<Option<Cow<'_, Record>>> {
30        ids.iter().map(|id| self.get(*id)).collect()
31    }
32
33    /// Total number of records held.
34    fn len(&self) -> usize;
35
36    fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39}
40
41// ── VecRecordStore (default in-memory impl) ───────────────────────────────────
42
43struct VecRecordStoreInner {
44    records:   Vec<Record>,
45    id_to_idx: HashMap<RecordId, usize>,
46}
47
48/// Default in-memory [`RecordStore`] backed by a `Vec`, zero-config.
49pub struct VecRecordStore {
50    inner: RwLock<VecRecordStoreInner>,
51}
52
53impl VecRecordStore {
54    pub fn new() -> Self {
55        Self {
56            inner: RwLock::new(VecRecordStoreInner {
57                records:   Vec::new(),
58                id_to_idx: HashMap::new(),
59            }),
60        }
61    }
62}
63
64impl Default for VecRecordStore {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl RecordStore for VecRecordStore {
71    fn insert(&self, record: Record) {
72        let mut inner = self.inner.write().unwrap();
73        let idx = inner.records.len();
74        inner.id_to_idx.insert(record.id, idx);
75        inner.records.push(record);
76    }
77
78    fn get(&self, id: RecordId) -> Option<Cow<'_, Record>> {
79        let inner = self.inner.read().unwrap();
80        let idx = *inner.id_to_idx.get(&id)?;
81        Some(Cow::Owned(inner.records[idx].clone()))
82    }
83
84    fn len(&self) -> usize {
85        self.inner.read().unwrap().records.len()
86    }
87}
88
89// ── BlockIndex ────────────────────────────────────────────────────────────────
90
91/// Opaque blocking index.
92///
93/// The `as_any` / `as_any_mut` escape hatches allow access to concrete fields
94/// not covered by the trait, such as index statistics.
95pub trait BlockIndex: Send + Sync {
96    /// Index `record_id` under the given set of blocking keys.
97    fn insert(&mut self, record_id: RecordId, keys: Vec<String>);
98
99    /// Return all record IDs sharing at least one key with `keys`, excluding
100    /// `exclude` (the querying record itself).  Result must be deduplicated.
101    fn lookup_union(&self, keys: &[String], exclude: RecordId) -> Vec<RecordId>;
102
103    /// Remove all index entries for `record_id`.
104    fn remove(&mut self, record_id: RecordId);
105
106    fn as_any(&self) -> &dyn std::any::Any;
107    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
108}
109
110/// Extracts blocking keys from records and looks up candidates in an index.
111pub trait Blocker: Send + Sync {
112    fn blocking_keys(&self, record: &Record, schema: &Schema) -> Vec<String>;
113    fn index_record(&self, record: &Record, schema: &Schema, index: &mut dyn BlockIndex);
114    fn candidates(&self, record: &Record, schema: &Schema, index: &dyn BlockIndex) -> Vec<RecordId>;
115}
116
117pub trait Comparator: Send + Sync {
118    /// Compare a single pair, always CPU, returns an individual vector.
119    fn compare(&self, a: &Record, b: &Record, schema: &Schema) -> ComparisonVector;
120
121    /// Pool-native batch comparison, the primary hot path.
122    ///
123    /// Reads `RecordPool` columns directly: zero HashMap lookups, no
124    /// `Record::clone()`.  Implementors SHOULD override this method.
125    /// The default falls back to `compare` per pair, which is correct but
126    /// slower than a native pool implementation.
127    fn compare_batch_from_pool(
128        &self,
129        pool:    &RecordPool,
130        indices: &[(usize, usize)],
131        schema:  &Schema,
132    ) -> ComparisonBatch {
133        let n_pairs  = indices.len();
134        let n_fields = schema.fields.len();
135        if n_pairs == 0 {
136            return ComparisonBatch::new(0, n_fields, vec![]);
137        }
138        let pair_ids: Vec<(u64, u64)> = indices.iter()
139            .map(|&(i, j)| (pool.ids[i], pool.ids[j]))
140            .collect();
141        let mut batch = ComparisonBatch::new(n_pairs, n_fields, pair_ids);
142        for (p, &(i, j)) in indices.iter().enumerate() {
143            use crate::record::FieldValue;
144            let mut a = Record::new(pool.ids[i]);
145            let mut b = Record::new(pool.ids[j]);
146            for (f, field) in schema.fields.iter().enumerate() {
147                let va = pool.get(f, i);
148                let vb = pool.get(f, j);
149                if !va.is_empty() {
150                    a = a.insert(&field.name, FieldValue::Text(va.to_string()));
151                }
152                if !vb.is_empty() {
153                    b = b.insert(&field.name, FieldValue::Text(vb.to_string()));
154                }
155            }
156            let v = self.compare(&a, &b, schema);
157            for (f, &level) in v.levels.iter().enumerate() {
158                batch.set_level(f, p, level);
159            }
160        }
161        batch
162    }
163
164}
165
166pub trait Scorer: Send + Sync {
167    /// Score a single pair, always CPU, cheap dot product.
168    fn score(&self, vector: &ComparisonVector, params: &ModelParams) -> ScoredPair;
169
170    /// Score a batch using the field-major `ComparisonBatch`.
171    fn score_batch(&self, batch: &ComparisonBatch, params: &ModelParams) -> Vec<ScoredPair> {
172        (0..batch.n_pairs)
173            .map(|p| self.score(&batch.pair_as_vector(p), params))
174            .collect()
175    }
176
177    fn estimate_params(
178        &self,
179        batch:    &ComparisonBatch,
180        init:     Option<ModelParams>,
181        max_iter: usize,
182    ) -> Result<ModelParams>;
183}
184
185/// Groups scored pairs into entity clusters.
186pub trait Clusterer: Send + Sync {
187    fn cluster(&self, pairs: &[ScoredPair], params: &ModelParams) -> Vec<Entity>;
188}
189
190/// Persistent store for resolved entities.
191pub trait EntityStore: Send + Sync {
192    fn upsert_entity(&self, entity: &Entity) -> Result<EntityId>;
193    fn get_entity(&self, id: EntityId) -> Result<Entity>;
194    fn record_to_entity(&self, record_id: RecordId) -> Result<Option<EntityId>>;
195    fn all_entities(&self) -> Result<Vec<Entity>>;
196}
197
198/// Convert an external row type into a [`Record`].
199///
200/// Implement this in an adapter crate (e.g. `zer-adapters`) for foreign
201/// row types such as a Polars `LazyFrame` row or an Arrow `RecordBatch` row.
202/// The `id` parameter lets callers assign a stable [`RecordId`].
203pub trait IntoRecord {
204    fn into_record(self, id: RecordId) -> Record;
205}
206
207impl IntoRecord for Record {
208    fn into_record(self, _id: RecordId) -> Record {
209        self
210    }
211}
212
213/// Neural re-ranker that adjudicates borderline record pairs.
214pub trait Judge: Send + Sync {
215    fn adjudicate(&self, pairs: &[ScoredPair]) -> Result<Vec<JudgeVerdict>>;
216}
217
218impl<J: Judge + ?Sized> Judge for Box<J> {
219    fn adjudicate(&self, pairs: &[ScoredPair]) -> Result<Vec<JudgeVerdict>> {
220        (**self).adjudicate(pairs)
221    }
222}
223
224#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
225pub enum JudgeVerdict {
226    IncreaseConfidence,
227    DecreaseConfidence,
228    NoChange,
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn traits_are_object_safe() {
237        let _: Box<dyn super::BlockIndex>;
238        let _: Box<dyn super::Blocker>;
239        let _: Box<dyn super::Comparator>;
240        let _: Box<dyn super::Scorer>;
241        let _: Box<dyn super::Clusterer>;
242        let _: Box<dyn super::EntityStore>;
243        let _: Box<dyn super::Judge>;
244        let _: Box<dyn super::RecordStore>;
245    }
246
247    #[test]
248    fn vec_record_store_insert_and_get() {
249        use crate::record::{FieldValue, Record};
250        let store = VecRecordStore::new();
251        assert!(store.is_empty());
252
253        let r = Record::new(42).insert("name", FieldValue::Text("Alice".into()));
254        store.insert(r);
255
256        assert_eq!(store.len(), 1);
257        let fetched = store.get(42).expect("record 42 must exist");
258        assert_eq!(fetched.id, 42);
259    }
260
261    #[test]
262    fn vec_record_store_get_missing_returns_none() {
263        let store = VecRecordStore::new();
264        assert!(store.get(999).is_none());
265    }
266
267    #[test]
268    fn vec_record_store_get_many() {
269        use crate::record::Record;
270        let store = VecRecordStore::new();
271        store.insert(Record::new(1));
272        store.insert(Record::new(2));
273        store.insert(Record::new(3));
274
275        let results = store.get_many(&[1, 3, 99]);
276        assert!(results[0].is_some());
277        assert!(results[1].is_some());
278        assert!(results[2].is_none());
279    }
280
281    #[test]
282    fn vec_record_store_is_sendable() {
283        use std::sync::Arc;
284        let store: Arc<dyn RecordStore> = Arc::new(VecRecordStore::new());
285        let store2 = Arc::clone(&store);
286        let handle = std::thread::spawn(move || {
287            store2.insert(crate::record::Record::new(7));
288        });
289        handle.join().unwrap();
290        assert_eq!(store.len(), 1);
291    }
292}
zer_core/traits.rs

zer_core/
traits.rs
