sbol_db_postgres/repo/
sketch.rs

1//! The MinHash/LSH similarity sketch store over Postgres.
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 single
8//! `band_hash = ANY(...)` posting-list union.
9
10use async_trait::async_trait;
11use sbol_db_core::DomainError;
12use sbol_db_storage::{Signature, SketchStore};
13use sqlx::Row;
14
15use crate::repo::db_err;
16use crate::PgPool;
17
18#[derive(Clone)]
19pub struct PgSketchStore {
20    pool: PgPool,
21}
22
23impl PgSketchStore {
24    pub fn new(pool: PgPool) -> Self {
25        Self { pool }
26    }
27}
28
29/// Reinterpret a 64-bit band hash as a signed `bigint` for storage. The search
30/// layer only tests band_hash equality, so the sign carries no meaning.
31fn to_i64(band: u64) -> i64 {
32    i64::from_le_bytes(band.to_le_bytes())
33}
34
35fn decode_sig(bytes: &[u8]) -> Result<Signature, DomainError> {
36    Signature::from_bytes(bytes).ok_or_else(|| {
37        DomainError::Database("stored sketch signature has an invalid length".into())
38    })
39}
40
41#[async_trait]
42impl SketchStore for PgSketchStore {
43    async fn put_sketch(
44        &self,
45        iri: &str,
46        signature: &Signature,
47        bands: &[u64],
48    ) -> Result<(), DomainError> {
49        let mut tx = self.pool.begin().await.map_err(db_err)?;
50        sqlx::query("DELETE FROM sbol_sequence_lsh_band WHERE sequence_iri = $1")
51            .bind(iri)
52            .execute(&mut *tx)
53            .await
54            .map_err(db_err)?;
55        sqlx::query(
56            r#"
57            INSERT INTO sbol_sequence_sketch (sequence_iri, signature)
58            VALUES ($1, $2)
59            ON CONFLICT (sequence_iri) DO UPDATE SET signature = EXCLUDED.signature
60            "#,
61        )
62        .bind(iri)
63        .bind(signature.to_bytes())
64        .execute(&mut *tx)
65        .await
66        .map_err(db_err)?;
67        if !bands.is_empty() {
68            let hashes: Vec<i64> = bands.iter().map(|b| to_i64(*b)).collect();
69            let iris: Vec<String> = vec![iri.to_owned(); bands.len()];
70            sqlx::query(
71                r#"
72                INSERT INTO sbol_sequence_lsh_band (band_hash, sequence_iri)
73                SELECT * FROM UNNEST($1::bigint[], $2::text[])
74                ON CONFLICT DO NOTHING
75                "#,
76            )
77            .bind(&hashes)
78            .bind(&iris)
79            .execute(&mut *tx)
80            .await
81            .map_err(db_err)?;
82        }
83        tx.commit().await.map_err(db_err)?;
84        Ok(())
85    }
86
87    async fn sketch_of(&self, iri: &str) -> Result<Option<Signature>, DomainError> {
88        let row = sqlx::query("SELECT signature FROM sbol_sequence_sketch WHERE sequence_iri = $1")
89            .bind(iri)
90            .fetch_optional(&self.pool)
91            .await
92            .map_err(db_err)?;
93        match row {
94            Some(r) => {
95                let bytes: Vec<u8> = r.try_get("signature").map_err(db_err)?;
96                Ok(Some(decode_sig(&bytes)?))
97            }
98            None => Ok(None),
99        }
100    }
101
102    async fn candidates_by_bands(&self, bands: &[u64]) -> Result<Vec<String>, DomainError> {
103        if bands.is_empty() {
104            return Ok(Vec::new());
105        }
106        let hashes: Vec<i64> = bands.iter().map(|b| to_i64(*b)).collect();
107        let rows = sqlx::query(
108            "SELECT DISTINCT sequence_iri FROM sbol_sequence_lsh_band WHERE band_hash = ANY($1::bigint[])",
109        )
110        .bind(&hashes)
111        .fetch_all(&self.pool)
112        .await
113        .map_err(db_err)?;
114        let mut out = Vec::with_capacity(rows.len());
115        for r in rows {
116            out.push(r.try_get::<String, _>("sequence_iri").map_err(db_err)?);
117        }
118        Ok(out)
119    }
120
121    async fn all_sketches(&self) -> Result<Vec<(String, Signature)>, DomainError> {
122        let rows = sqlx::query("SELECT sequence_iri, signature FROM sbol_sequence_sketch")
123            .fetch_all(&self.pool)
124            .await
125            .map_err(db_err)?;
126        let mut out = Vec::with_capacity(rows.len());
127        for r in rows {
128            let iri: String = r.try_get("sequence_iri").map_err(db_err)?;
129            let bytes: Vec<u8> = r.try_get("signature").map_err(db_err)?;
130            out.push((iri, decode_sig(&bytes)?));
131        }
132        Ok(out)
133    }
134
135    async fn replace_all_sketches(
136        &self,
137        entries: Vec<(String, Signature, Vec<u64>)>,
138    ) -> Result<(), DomainError> {
139        let mut tx = self.pool.begin().await.map_err(db_err)?;
140        sqlx::query("DELETE FROM sbol_sequence_lsh_band")
141            .execute(&mut *tx)
142            .await
143            .map_err(db_err)?;
144        sqlx::query("DELETE FROM sbol_sequence_sketch")
145            .execute(&mut *tx)
146            .await
147            .map_err(db_err)?;
148
149        let mut sketch_iris: Vec<String> = Vec::with_capacity(entries.len());
150        let mut sketch_sigs: Vec<Vec<u8>> = Vec::with_capacity(entries.len());
151        let mut band_hashes: Vec<i64> = Vec::new();
152        let mut band_iris: Vec<String> = Vec::new();
153        for (iri, signature, bands) in &entries {
154            sketch_iris.push(iri.clone());
155            sketch_sigs.push(signature.to_bytes());
156            for band in bands {
157                band_hashes.push(to_i64(*band));
158                band_iris.push(iri.clone());
159            }
160        }
161
162        if !sketch_iris.is_empty() {
163            sqlx::query(
164                r#"
165                INSERT INTO sbol_sequence_sketch (sequence_iri, signature)
166                SELECT * FROM UNNEST($1::text[], $2::bytea[])
167                "#,
168            )
169            .bind(&sketch_iris)
170            .bind(&sketch_sigs)
171            .execute(&mut *tx)
172            .await
173            .map_err(db_err)?;
174        }
175        if !band_hashes.is_empty() {
176            sqlx::query(
177                r#"
178                INSERT INTO sbol_sequence_lsh_band (band_hash, sequence_iri)
179                SELECT * FROM UNNEST($1::bigint[], $2::text[])
180                ON CONFLICT DO NOTHING
181                "#,
182            )
183            .bind(&band_hashes)
184            .bind(&band_iris)
185            .execute(&mut *tx)
186            .await
187            .map_err(db_err)?;
188        }
189        tx.commit().await.map_err(db_err)?;
190        Ok(())
191    }
192}
sbol_db_postgres/repo/sketch.rs

sbol_db_postgres/repo/
sketch.rs
