feat(vectors): add sparse vector foundation for playlist suggestions
Phase 1 of playlist-suggestions feature: - SparseVector: memory-efficient vector with math operations - set/get with sorted indices for O(log n) access - add, sub, scale, dot product, magnitude - normalize, cosine_similarity, top_k_similar - ArtistVectorStore: SQLite persistence - artist_index table: MBID to integer index mapping - vector_entries table: sparse representation (one row per non-zero) - vector_metadata table: track freshness with TTL - find_nearest: cosine similarity search with exclusion - 17 unit tests all passing
This commit is contained in:
vicrodh committed
1 parent
c1aaa87f8e
commit
1d2e28c1a8
5 files changed
+1016
No files matched your search
@@ -109,6 +109,9 @@ tauri-plugin-dialog = { version = "2.6.0" }
|
||||
[target.'cfg(target_os = "windows")'.dependencies]
|
||||
tauri-plugin-dialog = { version = "2.6.0" }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
|
||||
[patch.crates-io]
|
||||
cpal = { path = "vendor/cpal" }
|
||||
rodio = { path = "vendor/rodio" }
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Artist vector module for similarity-based discovery
|
||||
//!
|
||||
//! This module provides sparse vector representations of artists based on their
|
||||
//! relationships (from MusicBrainz) and similarity (from Qobuz). These vectors
|
||||
//! enable playlist suggestions by finding artists geometrically close to a
|
||||
//! playlist's combined vector.
|
||||
//!
|
||||
//! ## Architecture
|
||||
//!
|
||||
//! - `sparse_vector.rs`: Memory-efficient sparse vector with math operations
|
||||
//! - `store.rs`: SQLite persistence for artist indices and vectors
|
||||
//! - `weights.rs`: Configurable weights for different relationship types (Phase 2)
|
||||
//! - `builder.rs`: Vector construction from MusicBrainz + Qobuz (Phase 2)
|
||||
//! - `suggestions.rs`: Playlist suggestion engine (Phase 3)
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
//! ```ignore
|
||||
//! // Build vectors for artists
|
||||
//! let builder = ArtistVectorBuilder::new(store, mb_client, qobuz_client);
|
||||
//! builder.ensure_vector("artist-mbid", 7).await?;
|
||||
//!
|
||||
//! // Compute playlist vector
|
||||
//! let playlist_vec = compute_playlist_vector(&artist_mbids);
|
||||
//!
|
||||
//! // Find similar artists
|
||||
//! let suggestions = store.find_nearest(&playlist_vec, 10, &exclude)?;
|
||||
//! ```
|
||||
//!
|
||||
//! ## Vector Representation
|
||||
//!
|
||||
//! Each artist is represented as a sparse vector where:
|
||||
//! - Dimensions correspond to other artists (via integer index)
|
||||
//! - Values are weighted relationship strengths
|
||||
//! - Sources include: MusicBrainz relationships, Qobuz similarity
|
||||
//!
|
||||
//! Example:
|
||||
//! ```text
|
||||
//! Pink Floyd vector:
|
||||
//! [Roger Waters: 1.0 (member_of_band),
|
||||
//! David Gilmour: 1.0 (member_of_band),
|
||||
//! King Crimson: 0.7 (qobuz_similar),
|
||||
//! ...]
|
||||
//! ```
|
||||
|
||||
pub mod sparse_vector;
|
||||
pub mod store;
|
||||
|
||||
pub use sparse_vector::SparseVector;
|
||||
pub use store::{ArtistVectorStore, ArtistVectorStoreState, SimilarArtist, StoreStats, VectorEntry};
|
||||
@@ -0,0 +1,409 @@
|
||||
//! Sparse vector implementation for artist similarity
|
||||
//!
|
||||
//! Efficient representation for high-dimensional vectors where most values are zero.
|
||||
//! Used to represent artist relationships where each dimension is another artist.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// A sparse vector storing only non-zero values
|
||||
///
|
||||
/// Internally uses parallel vectors for indices and values for memory efficiency.
|
||||
/// For operations, temporarily converts to HashMap for O(1) lookups.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SparseVector {
|
||||
/// Sorted indices of non-zero elements
|
||||
indices: Vec<u32>,
|
||||
/// Values corresponding to indices
|
||||
values: Vec<f32>,
|
||||
}
|
||||
|
||||
impl SparseVector {
|
||||
/// Create an empty sparse vector
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
indices: Vec::new(),
|
||||
values: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a sparse vector with pre-allocated capacity
|
||||
pub fn with_capacity(capacity: usize) -> Self {
|
||||
Self {
|
||||
indices: Vec::with_capacity(capacity),
|
||||
values: Vec::with_capacity(capacity),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create from parallel vectors (must be same length, indices should be sorted)
|
||||
pub fn from_parts(indices: Vec<u32>, values: Vec<f32>) -> Self {
|
||||
debug_assert_eq!(indices.len(), values.len());
|
||||
Self { indices, values }
|
||||
}
|
||||
|
||||
/// Set a value at the given index
|
||||
///
|
||||
/// If the value is 0 (or very close), removes the entry.
|
||||
/// If the index exists, updates the value.
|
||||
/// If the index doesn't exist, inserts in sorted order.
|
||||
pub fn set(&mut self, idx: u32, value: f32) {
|
||||
// Ignore near-zero values
|
||||
if value.abs() < 1e-9 {
|
||||
self.remove(idx);
|
||||
return;
|
||||
}
|
||||
|
||||
match self.indices.binary_search(&idx) {
|
||||
Ok(pos) => {
|
||||
// Index exists, update value
|
||||
self.values[pos] = value;
|
||||
}
|
||||
Err(pos) => {
|
||||
// Index doesn't exist, insert at sorted position
|
||||
self.indices.insert(pos, idx);
|
||||
self.values.insert(pos, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the value at the given index (returns 0 if not present)
|
||||
pub fn get(&self, idx: u32) -> f32 {
|
||||
match self.indices.binary_search(&idx) {
|
||||
Ok(pos) => self.values[pos],
|
||||
Err(_) => 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove an entry by index
|
||||
pub fn remove(&mut self, idx: u32) {
|
||||
if let Ok(pos) = self.indices.binary_search(&idx) {
|
||||
self.indices.remove(pos);
|
||||
self.values.remove(pos);
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of non-zero elements
|
||||
pub fn nnz(&self) -> usize {
|
||||
self.indices.len()
|
||||
}
|
||||
|
||||
/// Check if vector is empty (all zeros)
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.indices.is_empty()
|
||||
}
|
||||
|
||||
/// Get indices of non-zero elements
|
||||
pub fn indices(&self) -> &[u32] {
|
||||
&self.indices
|
||||
}
|
||||
|
||||
/// Get values of non-zero elements
|
||||
pub fn values(&self) -> &[f32] {
|
||||
&self.values
|
||||
}
|
||||
|
||||
/// Iterate over (index, value) pairs
|
||||
pub fn iter(&self) -> impl Iterator<Item = (u32, f32)> + '_ {
|
||||
self.indices.iter().copied().zip(self.values.iter().copied())
|
||||
}
|
||||
|
||||
/// Add two sparse vectors element-wise
|
||||
pub fn add(&self, other: &SparseVector) -> SparseVector {
|
||||
let mut result = HashMap::new();
|
||||
|
||||
// Add all elements from self
|
||||
for (idx, val) in self.iter() {
|
||||
*result.entry(idx).or_insert(0.0) += val;
|
||||
}
|
||||
|
||||
// Add all elements from other
|
||||
for (idx, val) in other.iter() {
|
||||
*result.entry(idx).or_insert(0.0) += val;
|
||||
}
|
||||
|
||||
// Convert back to sparse vector
|
||||
let mut indices: Vec<u32> = result.keys().copied().collect();
|
||||
indices.sort_unstable();
|
||||
|
||||
let values: Vec<f32> = indices.iter().map(|idx| result[idx]).collect();
|
||||
|
||||
SparseVector { indices, values }
|
||||
}
|
||||
|
||||
/// Subtract another vector from this one (self - other)
|
||||
pub fn sub(&self, other: &SparseVector) -> SparseVector {
|
||||
let mut result = HashMap::new();
|
||||
|
||||
for (idx, val) in self.iter() {
|
||||
*result.entry(idx).or_insert(0.0) += val;
|
||||
}
|
||||
|
||||
for (idx, val) in other.iter() {
|
||||
*result.entry(idx).or_insert(0.0) -= val;
|
||||
}
|
||||
|
||||
let mut indices: Vec<u32> = result
|
||||
.iter()
|
||||
.filter(|(_, v)| v.abs() > 1e-9)
|
||||
.map(|(k, _)| *k)
|
||||
.collect();
|
||||
indices.sort_unstable();
|
||||
|
||||
let values: Vec<f32> = indices.iter().map(|idx| result[idx]).collect();
|
||||
|
||||
SparseVector { indices, values }
|
||||
}
|
||||
|
||||
/// Scale vector by a scalar
|
||||
pub fn scale(&self, scalar: f32) -> SparseVector {
|
||||
SparseVector {
|
||||
indices: self.indices.clone(),
|
||||
values: self.values.iter().map(|v| v * scalar).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute dot product with another sparse vector
|
||||
pub fn dot(&self, other: &SparseVector) -> f32 {
|
||||
let mut sum = 0.0;
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
|
||||
// Merge-style iteration (both vectors are sorted by index)
|
||||
while i < self.indices.len() && j < other.indices.len() {
|
||||
match self.indices[i].cmp(&other.indices[j]) {
|
||||
std::cmp::Ordering::Less => i += 1,
|
||||
std::cmp::Ordering::Greater => j += 1,
|
||||
std::cmp::Ordering::Equal => {
|
||||
sum += self.values[i] * other.values[j];
|
||||
i += 1;
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sum
|
||||
}
|
||||
|
||||
/// Compute L2 norm (magnitude) of the vector
|
||||
pub fn magnitude(&self) -> f32 {
|
||||
self.values.iter().map(|v| v * v).sum::<f32>().sqrt()
|
||||
}
|
||||
|
||||
/// Normalize the vector to unit length
|
||||
///
|
||||
/// Returns a zero vector if magnitude is zero.
|
||||
pub fn normalize(&self) -> SparseVector {
|
||||
let mag = self.magnitude();
|
||||
if mag < 1e-9 {
|
||||
return SparseVector::new();
|
||||
}
|
||||
|
||||
SparseVector {
|
||||
indices: self.indices.clone(),
|
||||
values: self.values.iter().map(|v| v / mag).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute cosine similarity with another vector
|
||||
///
|
||||
/// Returns 0 if either vector has zero magnitude.
|
||||
/// Result is in range [-1, 1], where 1 means identical direction.
|
||||
pub fn cosine_similarity(&self, other: &SparseVector) -> f32 {
|
||||
let dot = self.dot(other);
|
||||
let mag_self = self.magnitude();
|
||||
let mag_other = other.magnitude();
|
||||
|
||||
if mag_self < 1e-9 || mag_other < 1e-9 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
dot / (mag_self * mag_other)
|
||||
}
|
||||
|
||||
/// Find top-k most similar vectors from a collection
|
||||
///
|
||||
/// Returns vector of (index_in_collection, similarity) sorted by similarity descending.
|
||||
pub fn top_k_similar<'a>(
|
||||
&self,
|
||||
candidates: impl Iterator<Item = (usize, &'a SparseVector)>,
|
||||
k: usize,
|
||||
) -> Vec<(usize, f32)> {
|
||||
let mut scores: Vec<(usize, f32)> = candidates
|
||||
.map(|(idx, vec)| (idx, self.cosine_similarity(vec)))
|
||||
.filter(|(_, sim)| *sim > 0.0)
|
||||
.collect();
|
||||
|
||||
// Sort by similarity descending
|
||||
scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
scores.truncate(k);
|
||||
scores
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_set_and_get() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(5, 1.0);
|
||||
vec.set(10, 2.0);
|
||||
vec.set(3, 0.5);
|
||||
|
||||
assert_eq!(vec.get(5), 1.0);
|
||||
assert_eq!(vec.get(10), 2.0);
|
||||
assert_eq!(vec.get(3), 0.5);
|
||||
assert_eq!(vec.get(7), 0.0); // Not set
|
||||
assert_eq!(vec.nnz(), 3);
|
||||
|
||||
// Indices should be sorted
|
||||
assert_eq!(vec.indices(), &[3, 5, 10]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_value() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(5, 1.0);
|
||||
vec.set(5, 2.0);
|
||||
|
||||
assert_eq!(vec.get(5), 2.0);
|
||||
assert_eq!(vec.nnz(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_on_zero() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(5, 1.0);
|
||||
vec.set(5, 0.0);
|
||||
|
||||
assert_eq!(vec.get(5), 0.0);
|
||||
assert_eq!(vec.nnz(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add() {
|
||||
let mut a = SparseVector::new();
|
||||
a.set(1, 1.0);
|
||||
a.set(3, 2.0);
|
||||
|
||||
let mut b = SparseVector::new();
|
||||
b.set(2, 1.5);
|
||||
b.set(3, 0.5);
|
||||
|
||||
let c = a.add(&b);
|
||||
|
||||
assert_eq!(c.get(1), 1.0);
|
||||
assert_eq!(c.get(2), 1.5);
|
||||
assert_eq!(c.get(3), 2.5);
|
||||
assert_eq!(c.nnz(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dot_product() {
|
||||
let mut a = SparseVector::new();
|
||||
a.set(1, 2.0);
|
||||
a.set(3, 3.0);
|
||||
|
||||
let mut b = SparseVector::new();
|
||||
b.set(1, 4.0);
|
||||
b.set(2, 5.0);
|
||||
b.set(3, 1.0);
|
||||
|
||||
// dot = 2*4 + 3*1 = 11
|
||||
assert_eq!(a.dot(&b), 11.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_magnitude() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(0, 3.0);
|
||||
vec.set(1, 4.0);
|
||||
|
||||
// magnitude = sqrt(9 + 16) = 5
|
||||
assert!((vec.magnitude() - 5.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(0, 3.0);
|
||||
vec.set(1, 4.0);
|
||||
|
||||
let normalized = vec.normalize();
|
||||
|
||||
assert!((normalized.magnitude() - 1.0).abs() < 1e-6);
|
||||
assert!((normalized.get(0) - 0.6).abs() < 1e-6);
|
||||
assert!((normalized.get(1) - 0.8).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_similarity_identical() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(0, 1.0);
|
||||
vec.set(1, 2.0);
|
||||
|
||||
let sim = vec.cosine_similarity(&vec);
|
||||
assert!((sim - 1.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_similarity_orthogonal() {
|
||||
let mut a = SparseVector::new();
|
||||
a.set(0, 1.0);
|
||||
|
||||
let mut b = SparseVector::new();
|
||||
b.set(1, 1.0);
|
||||
|
||||
let sim = a.cosine_similarity(&b);
|
||||
assert!(sim.abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_similarity_similar() {
|
||||
let mut a = SparseVector::new();
|
||||
a.set(0, 1.0);
|
||||
a.set(1, 1.0);
|
||||
|
||||
let mut b = SparseVector::new();
|
||||
b.set(0, 1.0);
|
||||
b.set(1, 0.5);
|
||||
|
||||
let sim = a.cosine_similarity(&b);
|
||||
// Should be high but not 1.0
|
||||
assert!(sim > 0.9);
|
||||
assert!(sim < 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scale() {
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(0, 2.0);
|
||||
vec.set(1, 3.0);
|
||||
|
||||
let scaled = vec.scale(2.0);
|
||||
|
||||
assert_eq!(scaled.get(0), 4.0);
|
||||
assert_eq!(scaled.get(1), 6.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_parts() {
|
||||
let vec = SparseVector::from_parts(vec![1, 3, 5], vec![0.5, 1.0, 1.5]);
|
||||
|
||||
assert_eq!(vec.get(1), 0.5);
|
||||
assert_eq!(vec.get(3), 1.0);
|
||||
assert_eq!(vec.get(5), 1.5);
|
||||
assert_eq!(vec.nnz(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_vector() {
|
||||
let vec = SparseVector::new();
|
||||
|
||||
assert!(vec.is_empty());
|
||||
assert_eq!(vec.magnitude(), 0.0);
|
||||
assert_eq!(vec.normalize().nnz(), 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,553 @@
|
||||
//! SQLite-based storage for artist vectors
|
||||
//!
|
||||
//! Persists artist index mapping and sparse vectors for similarity search.
|
||||
|
||||
use rusqlite::{params, Connection, OptionalExtension};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::sparse_vector::SparseVector;
|
||||
|
||||
/// TTL for vector entries (7 days)
|
||||
const VECTOR_TTL_SECS: i64 = 7 * 24 * 60 * 60;
|
||||
|
||||
/// Shared state wrapper for Tauri
|
||||
pub struct ArtistVectorStoreState {
|
||||
pub store: Arc<Mutex<ArtistVectorStore>>,
|
||||
}
|
||||
|
||||
impl ArtistVectorStoreState {
|
||||
pub fn new() -> Result<Self, String> {
|
||||
let data_dir = dirs::data_dir()
|
||||
.ok_or("Could not determine data directory")?
|
||||
.join("qbz")
|
||||
.join("cache");
|
||||
|
||||
std::fs::create_dir_all(&data_dir)
|
||||
.map_err(|e| format!("Failed to create cache directory: {}", e))?;
|
||||
|
||||
let db_path = data_dir.join("artist_vectors.db");
|
||||
let store = ArtistVectorStore::new(&db_path)?;
|
||||
|
||||
log::info!("Artist vector store initialized at {:?}", db_path);
|
||||
|
||||
Ok(Self {
|
||||
store: Arc::new(Mutex::new(store)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Artist vector store with SQLite backend
|
||||
pub struct ArtistVectorStore {
|
||||
conn: Connection,
|
||||
/// In-memory cache of MBID to index mapping
|
||||
artist_to_idx: HashMap<String, u32>,
|
||||
/// Reverse mapping: index to MBID
|
||||
idx_to_artist: Vec<String>,
|
||||
/// Next available index
|
||||
next_idx: u32,
|
||||
}
|
||||
|
||||
/// Information about a stored vector entry
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct VectorEntry {
|
||||
pub artist_mbid: String,
|
||||
pub vector: SparseVector,
|
||||
pub updated_at: i64,
|
||||
pub sources: Vec<String>,
|
||||
}
|
||||
|
||||
/// Result of a similarity search
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SimilarArtist {
|
||||
pub mbid: String,
|
||||
pub name: Option<String>,
|
||||
pub similarity: f32,
|
||||
}
|
||||
|
||||
impl ArtistVectorStore {
|
||||
/// Create a new store at the given path
|
||||
pub fn new(path: &Path) -> Result<Self, String> {
|
||||
let conn = Connection::open(path)
|
||||
.map_err(|e| format!("Failed to open artist vector store: {}", e))?;
|
||||
|
||||
// Enable WAL mode for better concurrency
|
||||
conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL;")
|
||||
.map_err(|e| format!("Failed to set pragmas: {}", e))?;
|
||||
|
||||
let mut store = Self {
|
||||
conn,
|
||||
artist_to_idx: HashMap::new(),
|
||||
idx_to_artist: Vec::new(),
|
||||
next_idx: 0,
|
||||
};
|
||||
|
||||
store.init()?;
|
||||
store.load_artist_index()?;
|
||||
|
||||
Ok(store)
|
||||
}
|
||||
|
||||
/// Initialize database schema
|
||||
fn init(&self) -> Result<(), String> {
|
||||
self.conn
|
||||
.execute_batch(
|
||||
r#"
|
||||
-- Artist index: maps MBID to integer index for vectors
|
||||
CREATE TABLE IF NOT EXISTS artist_index (
|
||||
idx INTEGER PRIMARY KEY,
|
||||
mbid TEXT UNIQUE NOT NULL,
|
||||
name TEXT,
|
||||
created_at INTEGER NOT NULL DEFAULT (unixepoch())
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_artist_index_mbid ON artist_index(mbid);
|
||||
|
||||
-- Vector entries: sparse representation (one row per non-zero)
|
||||
CREATE TABLE IF NOT EXISTS vector_entries (
|
||||
artist_idx INTEGER NOT NULL,
|
||||
target_idx INTEGER NOT NULL,
|
||||
weight REAL NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
updated_at INTEGER NOT NULL DEFAULT (unixepoch()),
|
||||
PRIMARY KEY (artist_idx, target_idx, source)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_vector_entries_artist ON vector_entries(artist_idx);
|
||||
CREATE INDEX IF NOT EXISTS idx_vector_entries_target ON vector_entries(target_idx);
|
||||
CREATE INDEX IF NOT EXISTS idx_vector_entries_updated ON vector_entries(updated_at);
|
||||
|
||||
-- Vector metadata: track when each artist's vector was last computed
|
||||
CREATE TABLE IF NOT EXISTS vector_metadata (
|
||||
artist_idx INTEGER PRIMARY KEY,
|
||||
updated_at INTEGER NOT NULL,
|
||||
nnz INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
"#,
|
||||
)
|
||||
.map_err(|e| format!("Failed to initialize schema: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load artist index from database into memory
|
||||
fn load_artist_index(&mut self) -> Result<(), String> {
|
||||
let mut stmt = self
|
||||
.conn
|
||||
.prepare("SELECT idx, mbid FROM artist_index ORDER BY idx")
|
||||
.map_err(|e| format!("Failed to prepare index query: {}", e))?;
|
||||
|
||||
let rows = stmt
|
||||
.query_map([], |row| {
|
||||
Ok((row.get::<_, u32>(0)?, row.get::<_, String>(1)?))
|
||||
})
|
||||
.map_err(|e| format!("Failed to query index: {}", e))?;
|
||||
|
||||
self.artist_to_idx.clear();
|
||||
self.idx_to_artist.clear();
|
||||
|
||||
for row in rows {
|
||||
let (idx, mbid) = row.map_err(|e| format!("Failed to read row: {}", e))?;
|
||||
self.artist_to_idx.insert(mbid.clone(), idx);
|
||||
|
||||
// Ensure idx_to_artist has enough capacity
|
||||
while self.idx_to_artist.len() <= idx as usize {
|
||||
self.idx_to_artist.push(String::new());
|
||||
}
|
||||
self.idx_to_artist[idx as usize] = mbid;
|
||||
|
||||
if idx >= self.next_idx {
|
||||
self.next_idx = idx + 1;
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"Loaded {} artists into index, next_idx={}",
|
||||
self.artist_to_idx.len(),
|
||||
self.next_idx
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get or create an index for an artist MBID
|
||||
pub fn get_or_create_idx(&mut self, mbid: &str, name: Option<&str>) -> Result<u32, String> {
|
||||
if let Some(&idx) = self.artist_to_idx.get(mbid) {
|
||||
return Ok(idx);
|
||||
}
|
||||
|
||||
let idx = self.next_idx;
|
||||
self.next_idx += 1;
|
||||
|
||||
self.conn
|
||||
.execute(
|
||||
"INSERT INTO artist_index (idx, mbid, name) VALUES (?1, ?2, ?3)",
|
||||
params![idx, mbid, name],
|
||||
)
|
||||
.map_err(|e| format!("Failed to insert artist index: {}", e))?;
|
||||
|
||||
self.artist_to_idx.insert(mbid.to_string(), idx);
|
||||
|
||||
// Extend idx_to_artist
|
||||
while self.idx_to_artist.len() <= idx as usize {
|
||||
self.idx_to_artist.push(String::new());
|
||||
}
|
||||
self.idx_to_artist[idx as usize] = mbid.to_string();
|
||||
|
||||
Ok(idx)
|
||||
}
|
||||
|
||||
/// Get index for an artist MBID (returns None if not found)
|
||||
pub fn get_idx(&self, mbid: &str) -> Option<u32> {
|
||||
self.artist_to_idx.get(mbid).copied()
|
||||
}
|
||||
|
||||
/// Get MBID for an index
|
||||
pub fn get_mbid(&self, idx: u32) -> Option<&str> {
|
||||
self.idx_to_artist.get(idx as usize).map(|s| s.as_str())
|
||||
}
|
||||
|
||||
/// Store a vector for an artist
|
||||
pub fn set_vector(
|
||||
&mut self,
|
||||
mbid: &str,
|
||||
vector: &SparseVector,
|
||||
source: &str,
|
||||
) -> Result<(), String> {
|
||||
let artist_idx = self.get_or_create_idx(mbid, None)?;
|
||||
let now = current_timestamp();
|
||||
|
||||
// Delete existing entries for this artist+source
|
||||
self.conn
|
||||
.execute(
|
||||
"DELETE FROM vector_entries WHERE artist_idx = ?1 AND source = ?2",
|
||||
params![artist_idx, source],
|
||||
)
|
||||
.map_err(|e| format!("Failed to delete old entries: {}", e))?;
|
||||
|
||||
// Insert new entries
|
||||
let mut stmt = self
|
||||
.conn
|
||||
.prepare(
|
||||
"INSERT INTO vector_entries (artist_idx, target_idx, weight, source, updated_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)",
|
||||
)
|
||||
.map_err(|e| format!("Failed to prepare insert: {}", e))?;
|
||||
|
||||
for (target_idx, weight) in vector.iter() {
|
||||
stmt.execute(params![artist_idx, target_idx, weight, source, now])
|
||||
.map_err(|e| format!("Failed to insert entry: {}", e))?;
|
||||
}
|
||||
|
||||
// Update metadata
|
||||
self.conn
|
||||
.execute(
|
||||
"INSERT OR REPLACE INTO vector_metadata (artist_idx, updated_at, nnz)
|
||||
VALUES (?1, ?2, ?3)",
|
||||
params![artist_idx, now, vector.nnz()],
|
||||
)
|
||||
.map_err(|e| format!("Failed to update metadata: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the combined vector for an artist (all sources merged)
|
||||
pub fn get_vector(&self, mbid: &str) -> Option<SparseVector> {
|
||||
let artist_idx = self.get_idx(mbid)?;
|
||||
|
||||
let mut stmt = self
|
||||
.conn
|
||||
.prepare("SELECT target_idx, SUM(weight) FROM vector_entries WHERE artist_idx = ?1 GROUP BY target_idx")
|
||||
.ok()?;
|
||||
|
||||
let rows = stmt
|
||||
.query_map(params![artist_idx], |row| {
|
||||
Ok((row.get::<_, u32>(0)?, row.get::<_, f32>(1)?))
|
||||
})
|
||||
.ok()?;
|
||||
|
||||
let mut indices = Vec::new();
|
||||
let mut values = Vec::new();
|
||||
|
||||
for row in rows.flatten() {
|
||||
indices.push(row.0);
|
||||
values.push(row.1);
|
||||
}
|
||||
|
||||
if indices.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Sort by index
|
||||
let mut pairs: Vec<_> = indices.into_iter().zip(values).collect();
|
||||
pairs.sort_by_key(|(idx, _)| *idx);
|
||||
|
||||
let (indices, values): (Vec<_>, Vec<_>) = pairs.into_iter().unzip();
|
||||
|
||||
Some(SparseVector::from_parts(indices, values))
|
||||
}
|
||||
|
||||
/// Check if a vector exists and is fresh (within TTL)
|
||||
pub fn has_fresh_vector(&self, mbid: &str, max_age_secs: i64) -> bool {
|
||||
let Some(artist_idx) = self.get_idx(mbid) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let cutoff = current_timestamp() - max_age_secs;
|
||||
|
||||
let result: Option<i64> = self
|
||||
.conn
|
||||
.query_row(
|
||||
"SELECT updated_at FROM vector_metadata WHERE artist_idx = ?1 AND updated_at > ?2",
|
||||
params![artist_idx, cutoff],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.optional()
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
result.is_some()
|
||||
}
|
||||
|
||||
/// Find artists most similar to a query vector
|
||||
///
|
||||
/// Returns top-k artists by cosine similarity, excluding those in `exclude`.
|
||||
pub fn find_nearest(
|
||||
&self,
|
||||
query: &SparseVector,
|
||||
k: usize,
|
||||
exclude_mbids: &[String],
|
||||
) -> Result<Vec<SimilarArtist>, String> {
|
||||
if query.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Get all artist indices to exclude
|
||||
let exclude_indices: Vec<u32> = exclude_mbids
|
||||
.iter()
|
||||
.filter_map(|mbid| self.get_idx(mbid))
|
||||
.collect();
|
||||
|
||||
// Get all unique artist indices that have vectors
|
||||
let mut stmt = self
|
||||
.conn
|
||||
.prepare("SELECT DISTINCT artist_idx FROM vector_entries")
|
||||
.map_err(|e| format!("Failed to prepare query: {}", e))?;
|
||||
|
||||
let artist_indices: Vec<u32> = stmt
|
||||
.query_map([], |row| row.get(0))
|
||||
.map_err(|e| format!("Failed to query artists: {}", e))?
|
||||
.filter_map(|r| r.ok())
|
||||
.filter(|idx| !exclude_indices.contains(idx))
|
||||
.collect();
|
||||
|
||||
// Compute similarity for each candidate
|
||||
let mut similarities: Vec<(u32, f32)> = Vec::new();
|
||||
|
||||
for artist_idx in artist_indices {
|
||||
if let Some(mbid) = self.get_mbid(artist_idx) {
|
||||
if let Some(vec) = self.get_vector(mbid) {
|
||||
let sim = query.cosine_similarity(&vec);
|
||||
if sim > 0.0 {
|
||||
similarities.push((artist_idx, sim));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by similarity descending
|
||||
similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
similarities.truncate(k);
|
||||
|
||||
// Convert to result type
|
||||
let results = similarities
|
||||
.into_iter()
|
||||
.filter_map(|(idx, sim)| {
|
||||
let mbid = self.get_mbid(idx)?.to_string();
|
||||
let name = self.get_artist_name(&mbid);
|
||||
Some(SimilarArtist {
|
||||
mbid,
|
||||
name,
|
||||
similarity: sim,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Get artist name from index
|
||||
fn get_artist_name(&self, mbid: &str) -> Option<String> {
|
||||
self.conn
|
||||
.query_row(
|
||||
"SELECT name FROM artist_index WHERE mbid = ?1",
|
||||
params![mbid],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.optional()
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Get statistics about the store
|
||||
pub fn get_stats(&self) -> Result<StoreStats, String> {
|
||||
let artist_count = self.artist_to_idx.len();
|
||||
|
||||
let vector_count: i64 = self
|
||||
.conn
|
||||
.query_row("SELECT COUNT(DISTINCT artist_idx) FROM vector_entries", [], |row| {
|
||||
row.get(0)
|
||||
})
|
||||
.map_err(|e| format!("Failed to count vectors: {}", e))?;
|
||||
|
||||
let entry_count: i64 = self
|
||||
.conn
|
||||
.query_row("SELECT COUNT(*) FROM vector_entries", [], |row| row.get(0))
|
||||
.map_err(|e| format!("Failed to count entries: {}", e))?;
|
||||
|
||||
Ok(StoreStats {
|
||||
artist_count,
|
||||
vector_count: vector_count as usize,
|
||||
entry_count: entry_count as usize,
|
||||
})
|
||||
}
|
||||
|
||||
/// Clean up expired entries
|
||||
pub fn cleanup_expired(&mut self, max_age_secs: i64) -> Result<usize, String> {
|
||||
let cutoff = current_timestamp() - max_age_secs;
|
||||
|
||||
let deleted = self
|
||||
.conn
|
||||
.execute(
|
||||
"DELETE FROM vector_entries WHERE updated_at < ?1",
|
||||
params![cutoff],
|
||||
)
|
||||
.map_err(|e| format!("Failed to delete expired entries: {}", e))?;
|
||||
|
||||
// Also clean up metadata
|
||||
self.conn
|
||||
.execute(
|
||||
"DELETE FROM vector_metadata WHERE updated_at < ?1",
|
||||
params![cutoff],
|
||||
)
|
||||
.map_err(|e| format!("Failed to delete expired metadata: {}", e))?;
|
||||
|
||||
Ok(deleted)
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about the vector store
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoreStats {
|
||||
pub artist_count: usize,
|
||||
pub vector_count: usize,
|
||||
pub entry_count: usize,
|
||||
}
|
||||
|
||||
/// Get current Unix timestamp
|
||||
fn current_timestamp() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
fn create_test_store() -> ArtistVectorStore {
|
||||
let dir = tempdir().unwrap();
|
||||
let db_path = dir.path().join("test_vectors.db");
|
||||
ArtistVectorStore::new(&db_path).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_artist_index() {
|
||||
let mut store = create_test_store();
|
||||
|
||||
let idx1 = store.get_or_create_idx("mbid-1", Some("Artist 1")).unwrap();
|
||||
let idx2 = store.get_or_create_idx("mbid-2", Some("Artist 2")).unwrap();
|
||||
let idx1_again = store.get_or_create_idx("mbid-1", None).unwrap();
|
||||
|
||||
assert_eq!(idx1, 0);
|
||||
assert_eq!(idx2, 1);
|
||||
assert_eq!(idx1_again, idx1); // Should return same index
|
||||
|
||||
assert_eq!(store.get_mbid(idx1), Some("mbid-1"));
|
||||
assert_eq!(store.get_mbid(idx2), Some("mbid-2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_store_and_retrieve_vector() {
|
||||
let mut store = create_test_store();
|
||||
|
||||
// Create some artist indices first
|
||||
store.get_or_create_idx("target-1", None).unwrap();
|
||||
store.get_or_create_idx("target-2", None).unwrap();
|
||||
|
||||
let mut vec = SparseVector::new();
|
||||
vec.set(0, 1.0); // target-1
|
||||
vec.set(1, 0.5); // target-2
|
||||
|
||||
store.set_vector("artist-a", &vec, "test").unwrap();
|
||||
|
||||
let retrieved = store.get_vector("artist-a").unwrap();
|
||||
assert_eq!(retrieved.get(0), 1.0);
|
||||
assert_eq!(retrieved.get(1), 0.5);
|
||||
assert_eq!(retrieved.nnz(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_nearest() {
|
||||
let mut store = create_test_store();
|
||||
|
||||
// Create artists
|
||||
let idx_a = store.get_or_create_idx("artist-a", None).unwrap();
|
||||
let idx_b = store.get_or_create_idx("artist-b", None).unwrap();
|
||||
let idx_c = store.get_or_create_idx("artist-c", None).unwrap();
|
||||
|
||||
// Artist A: connected to B
|
||||
let mut vec_a = SparseVector::new();
|
||||
vec_a.set(idx_b, 1.0);
|
||||
store.set_vector("artist-a", &vec_a, "test").unwrap();
|
||||
|
||||
// Artist B: connected to A and C
|
||||
let mut vec_b = SparseVector::new();
|
||||
vec_b.set(idx_a, 1.0);
|
||||
vec_b.set(idx_c, 0.5);
|
||||
store.set_vector("artist-b", &vec_b, "test").unwrap();
|
||||
|
||||
// Artist C: connected to B
|
||||
let mut vec_c = SparseVector::new();
|
||||
vec_c.set(idx_b, 0.8);
|
||||
store.set_vector("artist-c", &vec_c, "test").unwrap();
|
||||
|
||||
// Query similar to A
|
||||
let query = vec_a.clone();
|
||||
let results = store.find_nearest(&query, 5, &["artist-a".to_string()]).unwrap();
|
||||
|
||||
// B should be most similar (both connect to each other)
|
||||
assert!(!results.is_empty());
|
||||
// C should also appear (both connect to B)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fresh_vector_check() {
|
||||
let mut store = create_test_store();
|
||||
|
||||
let vec = SparseVector::new();
|
||||
store.set_vector("artist-a", &vec, "test").unwrap();
|
||||
|
||||
// Should be fresh with 1 day TTL
|
||||
assert!(store.has_fresh_vector("artist-a", 86400));
|
||||
|
||||
// Should not be fresh with 0 second TTL
|
||||
assert!(!store.has_fresh_vector("artist-a", 0));
|
||||
|
||||
// Non-existent artist
|
||||
assert!(!store.has_fresh_vector("nonexistent", 86400));
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
pub mod api;
|
||||
pub mod api_cache;
|
||||
pub mod artist_vectors;
|
||||
pub mod audio;
|
||||
pub mod cache;
|
||||
pub mod cast;
|
||||
|
||||
Reference in new issue
Block a user