complete LSH
Ziyang Hu
1 year ago
parent
683b66c0a7
commit
ba98a5c137
@ -0,0 +1,374 @@
|
||||
/*
|
||||
* Copyright 2023, The Cozo Project Authors.
|
||||
*
|
||||
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
|
||||
* If a copy of the MPL was not distributed with this file,
|
||||
* You can obtain one at https://mozilla.org/MPL/2.0/.
|
||||
*/
|
||||
|
||||
// Some ideas are from https://github.com/schelterlabs/rust-minhash
|
||||
|
||||
use crate::data::expr::{eval_bytecode, Bytecode, eval_bytecode_pred};
|
||||
use crate::data::tuple::Tuple;
|
||||
use crate::fts::tokenizer::TextAnalyzer;
|
||||
use crate::fts::TokenizerConfig;
|
||||
use crate::runtime::relation::RelationHandle;
|
||||
use crate::runtime::transact::SessionTx;
|
||||
use crate::{DataValue, Expr, SourceSpan, Symbol};
|
||||
use miette::{bail, miette, Result};
|
||||
use quadrature::integrate;
|
||||
use rand::{thread_rng, RngCore};
|
||||
use rustc_hash::FxHashMap;
|
||||
use smartstring::{LazyCompact, SmartString};
|
||||
use std::cmp::min;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use twox_hash::XxHash32;
|
||||
|
||||
impl<'a> SessionTx<'a> {
|
||||
pub(crate) fn del_lsh_index_item(
|
||||
&mut self,
|
||||
tuple: &[DataValue],
|
||||
bytes: Option<Vec<u8>>,
|
||||
idx_handle: &RelationHandle,
|
||||
inv_idx_handle: &RelationHandle,
|
||||
manifest: &MinHashLshIndexManifest,
|
||||
) -> Result<()> {
|
||||
let bytes = match bytes {
|
||||
None => {
|
||||
if let Some(mut found) = inv_idx_handle.get_val_only(self, tuple)? {
|
||||
let inv_key = inv_idx_handle.encode_key_for_store(tuple, Default::default())?;
|
||||
self.store_tx.del(&inv_key)?;
|
||||
match found.pop() {
|
||||
Some(DataValue::Bytes(b)) => b,
|
||||
_ => unreachable!(),
|
||||
}
|
||||
} else {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Some(b) => b,
|
||||
};
|
||||
|
||||
let mut key = Vec::with_capacity(bytes.len() + 2);
|
||||
key.push(DataValue::Bot);
|
||||
key.push(DataValue::Bot);
|
||||
key.extend_from_slice(tuple);
|
||||
for (i, chunk) in bytes
|
||||
.chunks_exact(manifest.r * std::mem::size_of::<u32>())
|
||||
.enumerate()
|
||||
{
|
||||
key[0] = DataValue::from(i as i64);
|
||||
key[1] = DataValue::Bytes(chunk.to_vec());
|
||||
let key_bytes = idx_handle.encode_key_for_store(&key, Default::default())?;
|
||||
self.store_tx.del(&key_bytes)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
pub(crate) fn put_lsh_index_item(
|
||||
&mut self,
|
||||
tuple: &[DataValue],
|
||||
extractor: &[Bytecode],
|
||||
stack: &mut Vec<DataValue>,
|
||||
tokenizer: &TextAnalyzer,
|
||||
rel_handle: &RelationHandle,
|
||||
idx_handle: &RelationHandle,
|
||||
inv_idx_handle: &RelationHandle,
|
||||
manifest: &MinHashLshIndexManifest,
|
||||
hash_perms: &HashPermutations,
|
||||
) -> Result<()> {
|
||||
if let Some(mut found) =
|
||||
inv_idx_handle.get_val_only(self, &tuple[..rel_handle.metadata.keys.len()])?
|
||||
{
|
||||
let bytes = match found.pop() {
|
||||
Some(DataValue::Bytes(b)) => b,
|
||||
_ => unreachable!(),
|
||||
};
|
||||
self.del_lsh_index_item(tuple, Some(bytes), idx_handle, inv_idx_handle, manifest)?;
|
||||
}
|
||||
let to_index = eval_bytecode(extractor, tuple, stack)?;
|
||||
let min_hash = match to_index {
|
||||
DataValue::Null => return Ok(()),
|
||||
DataValue::List(l) => HashValues::new(l.iter(), hash_perms),
|
||||
DataValue::Str(s) => {
|
||||
let n_grams = tokenizer.unique_ngrams(&s, manifest.n_gram);
|
||||
HashValues::new(n_grams.iter(), hash_perms)
|
||||
}
|
||||
_ => bail!("Cannot put value {:?} into a LSH index", to_index),
|
||||
};
|
||||
let bytes = min_hash.get_bytes();
|
||||
let inv_key_part = &tuple[..rel_handle.metadata.keys.len()];
|
||||
let inv_val_part = vec![DataValue::Bytes(bytes.to_vec())];
|
||||
let inv_key = inv_idx_handle.encode_key_for_store(inv_key_part, Default::default())?;
|
||||
let inv_val =
|
||||
inv_idx_handle.encode_val_only_for_store(&inv_val_part, Default::default())?;
|
||||
self.store_tx.put(&inv_key, &inv_val)?;
|
||||
|
||||
let mut key = Vec::with_capacity(bytes.len() + 2);
|
||||
key.push(DataValue::Bot);
|
||||
key.push(DataValue::Bot);
|
||||
key.extend_from_slice(inv_key_part);
|
||||
let chunk_size = manifest.r * std::mem::size_of::<u32>();
|
||||
for i in 0..manifest.b {
|
||||
let byte_range = &bytes[i * chunk_size..(i + 1) * chunk_size];
|
||||
key[0] = DataValue::from(i as i64);
|
||||
key[1] = DataValue::Bytes(byte_range.to_vec());
|
||||
let key_bytes = idx_handle.encode_key_for_store(&key, Default::default())?;
|
||||
self.store_tx.put(&key_bytes, &[])?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
pub(crate) fn lsh_search(
|
||||
&self,
|
||||
tuple: &[DataValue],
|
||||
config: &LshSearch,
|
||||
stack: &mut Vec<DataValue>,
|
||||
filter_code: &Option<(Vec<Bytecode>, SourceSpan)>,
|
||||
) -> Result<Vec<Tuple>> {
|
||||
let bytes = if let Some(mut found) = config
|
||||
.inv_idx_handle
|
||||
.get_val_only(self, &tuple[..config.base_handle.metadata.keys.len()])?
|
||||
{
|
||||
match found.pop() {
|
||||
Some(DataValue::Bytes(b)) => b,
|
||||
_ => unreachable!(),
|
||||
}
|
||||
} else {
|
||||
return Ok(vec![]);
|
||||
};
|
||||
let chunk_size = config.manifest.r * std::mem::size_of::<u32>();
|
||||
let mut key_prefix = Vec::with_capacity(2);
|
||||
let mut found_tuples: FxHashMap<_, usize> = FxHashMap::default();
|
||||
for (i, chunk) in bytes.chunks_exact(chunk_size).enumerate() {
|
||||
key_prefix.clear();
|
||||
key_prefix.push(DataValue::from(i as i64));
|
||||
key_prefix.push(DataValue::Bytes(chunk.to_vec()));
|
||||
for ks in config.idx_handle.scan_prefix(self, &key_prefix) {
|
||||
let ks = ks?;
|
||||
let key_part = &ks[2..];
|
||||
if key_part != tuple {
|
||||
let found = found_tuples.entry(key_part.to_vec()).or_default();
|
||||
*found += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut ret = vec![];
|
||||
for (key, count) in found_tuples {
|
||||
let similarity = count as f64 / config.manifest.r as f64;
|
||||
if similarity < config.min_similarity {
|
||||
continue;
|
||||
}
|
||||
let mut orig_tuple = config
|
||||
.base_handle
|
||||
.get(self, &key)?
|
||||
.ok_or_else(|| miette!("Tuple not found in base LSH relation"))?;
|
||||
if let Some((filter_code, span)) = filter_code {
|
||||
if !eval_bytecode_pred(filter_code, &orig_tuple, stack, *span)? {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
if config.bind_similarity.is_some() {
|
||||
orig_tuple.push(DataValue::from(similarity));
|
||||
}
|
||||
ret.push(orig_tuple);
|
||||
if let Some(k) = config.k {
|
||||
if ret.len() >= k {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(ret)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct LshSearch {
|
||||
pub(crate) base_handle: RelationHandle,
|
||||
pub(crate) idx_handle: RelationHandle,
|
||||
pub(crate) inv_idx_handle: RelationHandle,
|
||||
pub(crate) manifest: MinHashLshIndexManifest,
|
||||
pub(crate) bindings: Vec<Symbol>,
|
||||
pub(crate) k: Option<usize>,
|
||||
pub(crate) query: Symbol,
|
||||
pub(crate) bind_similarity: Option<Symbol>,
|
||||
pub(crate) min_similarity: f64,
|
||||
// pub(crate) lax_mode: bool,
|
||||
pub(crate) filter: Option<Expr>,
|
||||
pub(crate) span: SourceSpan,
|
||||
}
|
||||
|
||||
impl LshSearch {
|
||||
pub(crate) fn all_bindings(&self) -> impl Iterator<Item = &Symbol> {
|
||||
self.bindings.iter().chain(self.bind_similarity.iter())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct HashValues(pub(crate) Vec<u32>);
|
||||
pub(crate) struct HashPermutations(pub(crate) Vec<u32>);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde_derive::Serialize, serde_derive::Deserialize)]
|
||||
pub(crate) struct MinHashLshIndexManifest {
|
||||
pub(crate) base_relation: SmartString<LazyCompact>,
|
||||
pub(crate) index_name: SmartString<LazyCompact>,
|
||||
pub(crate) extractor: String,
|
||||
pub(crate) n_gram: usize,
|
||||
pub(crate) tokenizer: TokenizerConfig,
|
||||
pub(crate) filters: Vec<TokenizerConfig>,
|
||||
|
||||
pub(crate) num_perm: usize,
|
||||
pub(crate) b: usize,
|
||||
pub(crate) r: usize,
|
||||
pub(crate) threshold: f64,
|
||||
pub(crate) perms: Vec<u8>,
|
||||
}
|
||||
|
||||
impl MinHashLshIndexManifest {
|
||||
pub(crate) fn get_hash_perms(&self) -> HashPermutations {
|
||||
HashPermutations::from_bytes(&self.perms)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct LshParams {
|
||||
pub b: usize,
|
||||
pub r: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Weights(pub(crate) f64, pub(crate) f64);
|
||||
|
||||
const _ALLOWED_INTEGRATE_ERR: f64 = 0.001;
|
||||
|
||||
// code is mostly from https://github.com/schelterlabs/rust-minhash/blob/81ea3fec24fd888a330a71b6932623643346b591/src/minhash_lsh.rs
|
||||
impl LshParams {
|
||||
pub fn find_optimal_params(threshold: f64, num_perm: usize, weights: &Weights) -> LshParams {
|
||||
let Weights(false_positive_weight, false_negative_weight) = weights;
|
||||
let mut min_error = f64::INFINITY;
|
||||
let mut opt = LshParams { b: 0, r: 0 };
|
||||
for b in 1..num_perm + 1 {
|
||||
let max_r = num_perm / b;
|
||||
for r in 1..max_r + 1 {
|
||||
let false_pos = LshParams::false_positive_probability(threshold, b, r);
|
||||
let false_neg = LshParams::false_negative_probability(threshold, b, r);
|
||||
let error = false_pos * false_positive_weight + false_neg * false_negative_weight;
|
||||
if error < min_error {
|
||||
min_error = error;
|
||||
opt = LshParams { b, r };
|
||||
}
|
||||
}
|
||||
}
|
||||
opt
|
||||
}
|
||||
|
||||
fn false_positive_probability(threshold: f64, b: usize, r: usize) -> f64 {
|
||||
let _probability =
|
||||
|s| -> f64 { 1. - f64::powf(1. - f64::powi(s, r as i32), b as f64) };
|
||||
integrate(_probability, 0.0, threshold, _ALLOWED_INTEGRATE_ERR).integral
|
||||
}
|
||||
|
||||
fn false_negative_probability(threshold: f64, b: usize, r: usize) -> f64 {
|
||||
let _probability =
|
||||
|s| -> f64 { 1. - (1. - f64::powf(1. - f64::powi(s, r as i32), b as f64)) };
|
||||
integrate(_probability, threshold, 1.0, _ALLOWED_INTEGRATE_ERR).integral
|
||||
}
|
||||
}
|
||||
|
||||
impl HashPermutations {
|
||||
pub(crate) fn new(n_perms: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
let mut perms = Vec::with_capacity(n_perms);
|
||||
for _ in 0..n_perms {
|
||||
perms.push(rng.next_u32());
|
||||
}
|
||||
Self(perms)
|
||||
}
|
||||
pub(crate) fn as_bytes(&self) -> &[u8] {
|
||||
unsafe {
|
||||
std::slice::from_raw_parts(
|
||||
self.0.as_ptr() as *const u8,
|
||||
self.0.len() * std::mem::size_of::<u32>(),
|
||||
)
|
||||
}
|
||||
}
|
||||
// this is the inverse of `as_bytes`
|
||||
pub(crate) fn from_bytes(bytes: &[u8]) -> Self {
|
||||
unsafe {
|
||||
let ptr = bytes.as_ptr() as *const u32;
|
||||
let len = bytes.len() / std::mem::size_of::<u32>();
|
||||
let perms = std::slice::from_raw_parts(ptr, len);
|
||||
Self(perms.to_vec())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HashValues {
|
||||
pub(crate) fn new<T: Hash>(values: impl Iterator<Item = T>, perms: &HashPermutations) -> Self {
|
||||
let mut ret = Self::init(perms);
|
||||
ret.update(values, perms);
|
||||
ret
|
||||
}
|
||||
pub(crate) fn init(perms: &HashPermutations) -> Self {
|
||||
Self(vec![u32::MAX; perms.0.len()])
|
||||
}
|
||||
pub(crate) fn update<T: Hash>(
|
||||
&mut self,
|
||||
values: impl Iterator<Item = T>,
|
||||
perms: &HashPermutations,
|
||||
) {
|
||||
for v in values {
|
||||
for (i, seed) in perms.0.iter().enumerate() {
|
||||
let mut hasher = XxHash32::with_seed(*seed);
|
||||
v.hash(&mut hasher);
|
||||
let hash = hasher.finish() as u32;
|
||||
self.0[i] = min(self.0[i], hash);
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
pub(crate) fn jaccard(&self, other_minhash: &Self) -> f32 {
|
||||
let matches = self
|
||||
.0
|
||||
.iter()
|
||||
.zip_eq(&other_minhash.0)
|
||||
.filter(|(left, right)| left == right)
|
||||
.count();
|
||||
let result = matches as f32 / self.0.len() as f32;
|
||||
result
|
||||
}
|
||||
pub(crate) fn get_bytes(&self) -> &[u8] {
|
||||
unsafe {
|
||||
std::slice::from_raw_parts(
|
||||
self.0.as_ptr() as *const u8,
|
||||
self.0.len() * std::mem::size_of::<u32>(),
|
||||
)
|
||||
}
|
||||
}
|
||||
// pub(crate) fn get_byte_chunks(&self, n_chunks: usize) -> impl Iterator<Item = &[u8]> {
|
||||
// let chunk_size = self.0.len() * std::mem::size_of::<u32>() / n_chunks;
|
||||
// self.get_bytes().chunks_exact(chunk_size)
|
||||
// }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_minhash() {
|
||||
let perms = HashPermutations::new(20000);
|
||||
let mut m1 = HashValues::new([1, 2, 3, 4, 5, 6].iter(), &perms);
|
||||
let mut m2 = HashValues::new([4, 3, 2, 1, 5, 6].iter(), &perms);
|
||||
assert_eq!(m1.0, m2.0);
|
||||
// println!("{:?}", &m1.0);
|
||||
// println!("{:?}", &m2.0);
|
||||
assert_eq!(m1.jaccard(&m2), 1.0);
|
||||
m1.update([7, 8, 9].iter(), &perms);
|
||||
assert!(m1.jaccard(&m2) < 1.0);
|
||||
println!("{:?}", m1.jaccard(&m2));
|
||||
m2.update([17, 18, 19].iter(), &perms);
|
||||
assert!(m1.jaccard(&m2) < 1.0);
|
||||
println!("{:?}", m1.jaccard(&m2));
|
||||
// println!("{:?}", m2.get_byte_chunks(2).collect_vec());
|
||||
assert_eq!(perms.0, HashPermutations::from_bytes(perms.as_bytes()).0);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue