sbol_db_sqlite/repo/
sketch.rs

1//! The MinHash/LSH similarity sketch store over SQLite.
2//!
3//! Signatures live in `sbol_sequence_sketch` (one BLOB per sequence) and the
4//! LSH band buckets in `sbol_sequence_lsh_band` (one row per band a sequence
5//! falls into). A per-sequence write replaces both in one transaction, so a
6//! re-index leaves no stale postings. Candidate generation
7//! ([`candidates_by_bands`](SketchStore::candidates_by_bands)) is a posting-list
8//! union over `band_hash`.
9
10use async_trait::async_trait;
11use sbol_db_core::DomainError;
12use sbol_db_storage::{Signature, SketchStore};
13use sqlx::Row;
14
15use crate::pool::db_err;
16use crate::SqlitePool;
17
18#[derive(Clone)]
19pub struct SqliteSketchStore {
20    pool: SqlitePool,
21}
22
23impl SqliteSketchStore {
24    pub fn new(pool: SqlitePool) -> Self {
25        Self { pool }
26    }
27}
28
29/// Reinterpret a 64-bit band hash as a signed integer for the INTEGER column.
30/// The search layer only tests band_hash equality, so the sign carries no
31/// meaning.
32fn to_i64(band: u64) -> i64 {
33    i64::from_le_bytes(band.to_le_bytes())
34}
35
36fn decode_sig(bytes: &[u8]) -> Result<Signature, DomainError> {
37    Signature::from_bytes(bytes).ok_or_else(|| {
38        DomainError::Database("stored sketch signature has an invalid length".into())
39    })
40}
41
42#[async_trait]
43impl SketchStore for SqliteSketchStore {
44    async fn put_sketch(
45        &self,
46        iri: &str,
47        signature: &Signature,
48        bands: &[u64],
49    ) -> Result<(), DomainError> {
50        let mut tx = self.pool.begin().await.map_err(db_err)?;
51        sqlx::query("DELETE FROM sbol_sequence_lsh_band WHERE sequence_iri = ?")
52            .bind(iri)
53            .execute(&mut *tx)
54            .await
55            .map_err(db_err)?;
56        sqlx::query(
57            "INSERT INTO sbol_sequence_sketch (sequence_iri, signature) VALUES (?, ?) \
58             ON CONFLICT (sequence_iri) DO UPDATE SET signature = excluded.signature",
59        )
60        .bind(iri)
61        .bind(signature.to_bytes())
62        .execute(&mut *tx)
63        .await
64        .map_err(db_err)?;
65        for band in bands {
66            sqlx::query(
67                "INSERT INTO sbol_sequence_lsh_band (band_hash, sequence_iri) VALUES (?, ?) \
68                 ON CONFLICT DO NOTHING",
69            )
70            .bind(to_i64(*band))
71            .bind(iri)
72            .execute(&mut *tx)
73            .await
74            .map_err(db_err)?;
75        }
76        tx.commit().await.map_err(db_err)?;
77        Ok(())
78    }
79
80    async fn sketch_of(&self, iri: &str) -> Result<Option<Signature>, DomainError> {
81        let row = sqlx::query("SELECT signature FROM sbol_sequence_sketch WHERE sequence_iri = ?")
82            .bind(iri)
83            .fetch_optional(&self.pool)
84            .await
85            .map_err(db_err)?;
86        match row {
87            Some(r) => {
88                let bytes: Vec<u8> = r.try_get("signature").map_err(db_err)?;
89                Ok(Some(decode_sig(&bytes)?))
90            }
91            None => Ok(None),
92        }
93    }
94
95    async fn candidates_by_bands(&self, bands: &[u64]) -> Result<Vec<String>, DomainError> {
96        if bands.is_empty() {
97            return Ok(Vec::new());
98        }
99        let placeholders = vec!["?"; bands.len()].join(", ");
100        let sql = format!(
101            "SELECT DISTINCT sequence_iri FROM sbol_sequence_lsh_band \
102             WHERE band_hash IN ({placeholders})"
103        );
104        let mut q = sqlx::query(&sql);
105        for band in bands {
106            q = q.bind(to_i64(*band));
107        }
108        let rows = q.fetch_all(&self.pool).await.map_err(db_err)?;
109        let mut out = Vec::with_capacity(rows.len());
110        for r in rows {
111            out.push(r.try_get::<String, _>("sequence_iri").map_err(db_err)?);
112        }
113        Ok(out)
114    }
115
116    async fn all_sketches(&self) -> Result<Vec<(String, Signature)>, DomainError> {
117        let rows = sqlx::query("SELECT sequence_iri, signature FROM sbol_sequence_sketch")
118            .fetch_all(&self.pool)
119            .await
120            .map_err(db_err)?;
121        let mut out = Vec::with_capacity(rows.len());
122        for r in rows {
123            let iri: String = r.try_get("sequence_iri").map_err(db_err)?;
124            let bytes: Vec<u8> = r.try_get("signature").map_err(db_err)?;
125            out.push((iri, decode_sig(&bytes)?));
126        }
127        Ok(out)
128    }
129
130    async fn replace_all_sketches(
131        &self,
132        entries: Vec<(String, Signature, Vec<u64>)>,
133    ) -> Result<(), DomainError> {
134        let mut tx = self.pool.begin().await.map_err(db_err)?;
135        sqlx::query("DELETE FROM sbol_sequence_lsh_band")
136            .execute(&mut *tx)
137            .await
138            .map_err(db_err)?;
139        sqlx::query("DELETE FROM sbol_sequence_sketch")
140            .execute(&mut *tx)
141            .await
142            .map_err(db_err)?;
143        for (iri, signature, bands) in &entries {
144            sqlx::query("INSERT INTO sbol_sequence_sketch (sequence_iri, signature) VALUES (?, ?)")
145                .bind(iri)
146                .bind(signature.to_bytes())
147                .execute(&mut *tx)
148                .await
149                .map_err(db_err)?;
150            for band in bands {
151                sqlx::query(
152                    "INSERT INTO sbol_sequence_lsh_band (band_hash, sequence_iri) VALUES (?, ?) \
153                     ON CONFLICT DO NOTHING",
154                )
155                .bind(to_i64(*band))
156                .bind(iri)
157                .execute(&mut *tx)
158                .await
159                .map_err(db_err)?;
160            }
161        }
162        tx.commit().await.map_err(db_err)?;
163        Ok(())
164    }
165}
sbol_db_sqlite/repo/sketch.rs

sbol_db_sqlite/repo/
sketch.rs
