From 2363156e3a8d99130f4dbcef4c81c2549fa98442 Mon Sep 17 00:00:00 2001 From: Timothy Quilling Date: Sat, 15 Nov 2025 23:47:09 -0500 Subject: [PATCH] feat: token based search --- Cargo.lock | 7 + consumer/Cargo.toml | 1 + .../src/database_writer/operations/mod.rs | 2 + .../operations/search_index.rs | 69 ++++++++ consumer/src/lib.rs | 1 + consumer/src/search/mod.rs | 2 + consumer/src/search/tokenizer.rs | 154 ++++++++++++++++++ .../down.sql | 1 + .../up.sql | 31 ++++ parakeet-db/src/models.rs | 27 +++ parakeet-db/src/schema.rs | 14 ++ 11 files changed, 309 insertions(+) create mode 100644 consumer/src/database_writer/operations/search_index.rs create mode 100644 consumer/src/search/mod.rs create mode 100644 consumer/src/search/tokenizer.rs create mode 100644 migrations/2025-11-15-225000_create_post_search_tokens/down.sql create mode 100644 migrations/2025-11-15-225000_create_post_search_tokens/up.sql diff --git a/Cargo.lock b/Cargo.lock index c7b5a197..14dbd793 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -684,6 +684,7 @@ dependencies = [ "tokio-util", "tracing", "tracing-subscriber", + "unicode-segmentation", "urlencoding", "zstd", ] @@ -4519,6 +4520,12 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e70f2a8b45122e719eb623c01822704c4e0907e7e426a05927e1a1cfff5b75d0" +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + [[package]] name = "unsigned-varint" version = "0.7.2" diff --git a/consumer/Cargo.toml b/consumer/Cargo.toml index 19ccdf7c..0f21588b 100644 --- a/consumer/Cargo.toml +++ b/consumer/Cargo.toml @@ -46,6 +46,7 @@ once_cell = "1.18.0" futures-channel = "0.3.31" urlencoding = "2.1" color-eyre = "0.6.5" +unicode-segmentation = "1.10" [lints.rust] diff --git a/consumer/src/database_writer/operations/mod.rs b/consumer/src/database_writer/operations/mod.rs index d56ec818..6e028b40 100644 --- a/consumer/src/database_writer/operations/mod.rs +++ b/consumer/src/database_writer/operations/mod.rs @@ -119,9 +119,11 @@ pub mod executor; pub mod handlers; pub mod notifications; pub mod processor; +pub mod search_index; pub mod types; pub use processor::process_record_to_operations; +pub use search_index::SearchIndexWriter; pub use types::{AggregateDelta, DatabaseOperation}; // SelfLabels available via types::SelfLabels if needed, but most code imports from lexica directly // Notification helpers re-exported for convenience diff --git a/consumer/src/database_writer/operations/search_index.rs b/consumer/src/database_writer/operations/search_index.rs new file mode 100644 index 00000000..83ef0b83 --- /dev/null +++ b/consumer/src/database_writer/operations/search_index.rs @@ -0,0 +1,69 @@ +use crate::search::SearchTokenizer; +use chrono::Utc; +use deadpool_postgres::GenericClient; +use eyre::Context as _; + +pub struct SearchIndexWriter { + tokenizer: SearchTokenizer, +} + +impl SearchIndexWriter { + pub fn new() -> Self { + Self { + tokenizer: SearchTokenizer::new(), + } + } + + /// Index posts for search + /// + /// Takes post data (id, content, author_did, etc.) and creates tokenized search records. + /// This is designed to be called from the database writer after posts are inserted. + pub async fn index_posts_raw( + &self, + conn: &C, + post_data: &[(i64, Option>, String, Option, bool, bool)], // (post_id, content, author_did, lang, has_media, has_links) + ) -> eyre::Result<()> { + if post_data.is_empty() { + return Ok(()); + } + + let now = Utc::now(); + + for (post_id, content, author_did, lang, has_media, has_links) in post_data { + // Decompress and tokenize content + let tokens = if let Some(content_bytes) = content { + let decompressed = zstd::decode_all(content_bytes.as_slice()) + .wrap_err("Failed to decompress post content")?; + let text = String::from_utf8(decompressed) + .wrap_err("Failed to decode post content as UTF-8")?; + self.tokenizer.tokenize(&text) + } else { + // Empty post or stub - no tokens + vec![] + }; + + // Insert or update search tokens + conn.execute( + "INSERT INTO post_search_tokens (post_id, indexed_at, tokens, lang, has_media, has_links, author_did) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (post_id) DO UPDATE SET + tokens = EXCLUDED.tokens, + indexed_at = EXCLUDED.indexed_at, + lang = EXCLUDED.lang, + has_media = EXCLUDED.has_media, + has_links = EXCLUDED.has_links", + &[post_id, &now, &tokens, lang, has_media, has_links, author_did], + ) + .await + .wrap_err("Failed to insert/update search tokens")?; + } + + Ok(()) + } +} + +impl Default for SearchIndexWriter { + fn default() -> Self { + Self::new() + } +} diff --git a/consumer/src/lib.rs b/consumer/src/lib.rs index 74d92c03..40149432 100644 --- a/consumer/src/lib.rs +++ b/consumer/src/lib.rs @@ -15,6 +15,7 @@ mod indexer; // mod label_indexer; // Disabled - will be reimplemented with new cursor system pub mod parsing; mod relay; +pub mod search; mod sources; pub mod types; mod utils; diff --git a/consumer/src/search/mod.rs b/consumer/src/search/mod.rs new file mode 100644 index 00000000..59b78555 --- /dev/null +++ b/consumer/src/search/mod.rs @@ -0,0 +1,2 @@ +pub mod tokenizer; +pub use tokenizer::SearchTokenizer; diff --git a/consumer/src/search/tokenizer.rs b/consumer/src/search/tokenizer.rs new file mode 100644 index 00000000..55c72f2b --- /dev/null +++ b/consumer/src/search/tokenizer.rs @@ -0,0 +1,154 @@ +use regex::Regex; +use std::collections::HashSet; +use std::sync::LazyLock; +use unicode_segmentation::UnicodeSegmentation; + +static HASHTAG_RE: LazyLock = LazyLock::new(|| Regex::new(r"#([\p{L}\p{N}_]+)").unwrap()); +static MENTION_RE: LazyLock = LazyLock::new(|| Regex::new(r"@([\w-]+)\.[\w.-]+").unwrap()); + +pub struct SearchTokenizer { + min_token_length: usize, + max_tokens_per_post: usize, + stopwords: HashSet, +} + +impl SearchTokenizer { + pub fn new() -> Self { + Self { + min_token_length: 2, + max_tokens_per_post: 100, + stopwords: Self::load_stopwords(), + } + } + + pub fn tokenize(&self, content: &str) -> Vec { + let mut tokens = Vec::new(); + let mut seen = HashSet::new(); + + // Extract hashtags (preserve as single tokens) + for cap in HASHTAG_RE.captures_iter(content) { + let tag = cap[1].to_lowercase(); + if tag.len() >= self.min_token_length && seen.insert(tag.clone()) { + tokens.push(tag); + } + } + + // Extract mentions (username only) + for cap in MENTION_RE.captures_iter(content) { + let username = cap[1].to_lowercase(); + if username.len() >= self.min_token_length && seen.insert(username.clone()) { + tokens.push(username); + } + } + + // Extract regular words + for word in UnicodeSegmentation::unicode_words(content) { + // Skip URLs + if word.starts_with("http") || word.contains(".com") || word.contains(".org") { + continue; + } + + let normalized = word.to_lowercase(); + + // Filter: too short, stopwords, already seen + if normalized.len() < self.min_token_length { + continue; + } + if self.stopwords.contains(&normalized) { + continue; + } + if !seen.insert(normalized.clone()) { + continue; + } + + tokens.push(normalized); + + if tokens.len() >= self.max_tokens_per_post { + break; + } + } + + tokens + } + + fn load_stopwords() -> HashSet { + // Top 100 English stopwords + let words = vec![ + "a", "an", "and", "are", "as", "at", "be", "by", "for", + "from", "has", "he", "in", "is", "it", "its", "of", "on", + "that", "the", "to", "was", "will", "with", "this", "but", + "they", "have", "had", "what", "when", "where", "who", "which", + "not", "or", "if", "so", "can", "been", "would", "could", + "should", "may", "do", "does", "did", "me", "we", "us", + "our", "your", "their", "his", "her", "she", "him", "them", + "my", "you", "i", "am", "up", "out", "about", "into", + "than", "then", "now", "all", "some", "any", "no", "more", + "most", "just", "only", "very", "too", "also", "such", + "much", "many", "even", "still", "how", "why", "here", + "there", "said", "each", "which", "these", "those", "other", + "another", "both", "few", "own", "same", "get", "got", + ]; + words.into_iter().map(String::from).collect() + } +} + +impl Default for SearchTokenizer { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_hashtag_extraction() { + let tokenizer = SearchTokenizer::new(); + let tokens = tokenizer.tokenize("Love #Rust and #Bluesky! #AI2024"); + + assert!(tokens.contains(&"rust".to_string())); + assert!(tokens.contains(&"bluesky".to_string())); + assert!(tokens.contains(&"ai2024".to_string())); + } + + #[test] + fn test_mention_extraction() { + let tokenizer = SearchTokenizer::new(); + let tokens = tokenizer.tokenize("Hey @alice.bsky.social check this!"); + + assert!(tokens.contains(&"alice".to_string())); + assert!(tokens.contains(&"hey".to_string())); + assert!(tokens.contains(&"check".to_string())); + } + + #[test] + fn test_stopword_filtering() { + let tokenizer = SearchTokenizer::new(); + let tokens = tokenizer.tokenize("The quick brown fox"); + + assert!(!tokens.contains(&"the".to_string())); + assert!(tokens.contains(&"quick".to_string())); + assert!(tokens.contains(&"brown".to_string())); + assert!(tokens.contains(&"fox".to_string())); + } + + #[test] + fn test_url_skipping() { + let tokenizer = SearchTokenizer::new(); + let tokens = tokenizer.tokenize("Check out https://example.com/article"); + + assert!(tokens.contains(&"check".to_string())); + assert!(!tokens.iter().any(|t| t.contains("http"))); + assert!(!tokens.iter().any(|t| t.contains("example"))); + } + + #[test] + fn test_deduplication() { + let tokenizer = SearchTokenizer::new(); + let tokens = tokenizer.tokenize("rust rust rust #Rust"); + + // Should only appear once (deduplicated) + assert_eq!(tokens.iter().filter(|t| *t == "rust").count(), 1); + } +} diff --git a/migrations/2025-11-15-225000_create_post_search_tokens/down.sql b/migrations/2025-11-15-225000_create_post_search_tokens/down.sql new file mode 100644 index 00000000..b11522f8 --- /dev/null +++ b/migrations/2025-11-15-225000_create_post_search_tokens/down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS post_search_tokens CASCADE; diff --git a/migrations/2025-11-15-225000_create_post_search_tokens/up.sql b/migrations/2025-11-15-225000_create_post_search_tokens/up.sql new file mode 100644 index 00000000..be955891 --- /dev/null +++ b/migrations/2025-11-15-225000_create_post_search_tokens/up.sql @@ -0,0 +1,31 @@ +-- Create search tokens table (regular PostgreSQL table) +-- TimescaleDB conversion will happen in a future migration +CREATE TABLE post_search_tokens ( + post_id BIGINT PRIMARY KEY, + indexed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + tokens TEXT[] NOT NULL, + lang TEXT, + has_media BOOLEAN NOT NULL DEFAULT false, + has_links BOOLEAN NOT NULL DEFAULT false, + author_did TEXT NOT NULL, + + CONSTRAINT fk_post FOREIGN KEY (post_id) + REFERENCES posts(id) ON DELETE CASCADE +); + +-- Indexes +CREATE INDEX idx_post_search_tokens_gin + ON post_search_tokens USING GIN (tokens); + +CREATE INDEX idx_post_search_author + ON post_search_tokens (author_did, indexed_at DESC); + +CREATE INDEX idx_post_search_lang + ON post_search_tokens (lang, indexed_at DESC) + WHERE lang IS NOT NULL; + +CREATE INDEX idx_post_search_indexed_at + ON post_search_tokens (indexed_at DESC); + +-- Note: We'll convert this to a TimescaleDB hypertable in a future migration +-- See .claude/plans/1115-2249-TOKEN_SEARCH_IMPLEMENTATION.md appendix for TimescaleDB conversion plan diff --git a/parakeet-db/src/models.rs b/parakeet-db/src/models.rs index 15bdfbbe..f6d1ad41 100644 --- a/parakeet-db/src/models.rs +++ b/parakeet-db/src/models.rs @@ -235,6 +235,33 @@ pub struct PostMention { pub mentioned_actor_id: i32, // PK: FK to actors } +// Post search tokens - for token-based full-text search +#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[diesel(table_name = crate::schema::post_search_tokens)] +#[diesel(primary_key(post_id))] +#[diesel(check_for_backend(diesel::pg::Pg))] +pub struct PostSearchToken { + pub post_id: i64, // PK: FK to posts + pub indexed_at: DateTime, + pub tokens: Vec>, // Token array for GIN index search + pub lang: Option, + pub has_media: bool, + pub has_links: bool, + pub author_did: String, +} + +#[derive(Debug, Insertable)] +#[diesel(table_name = crate::schema::post_search_tokens)] +pub struct NewPostSearchToken { + pub post_id: i64, + pub indexed_at: DateTime, + pub tokens: Vec, // Non-nullable tokens for insertion + pub lang: Option, + pub has_media: bool, + pub has_links: bool, + pub author_did: String, +} + #[derive(Debug, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::uris)] #[diesel(primary_key(id))] diff --git a/parakeet-db/src/schema.rs b/parakeet-db/src/schema.rs index 00572066..2066e55d 100644 --- a/parakeet-db/src/schema.rs +++ b/parakeet-db/src/schema.rs @@ -539,6 +539,18 @@ diesel::table! { } } +diesel::table! { + post_search_tokens (post_id) { + post_id -> Int8, + indexed_at -> Timestamptz, + tokens -> Array>, + lang -> Nullable, + has_media -> Bool, + has_links -> Bool, + author_did -> Text, + } +} + diesel::table! { postgate_detached (postgate_id, detached_post_id) { postgate_id -> Int8, @@ -760,6 +772,7 @@ diesel::joinable!(post_facets -> posts (post_id)); diesel::joinable!(post_facets -> uris (link_uri_id)); diesel::joinable!(post_mentions -> actors (mentioned_actor_id)); diesel::joinable!(post_mentions -> posts (post_id)); +diesel::joinable!(post_search_tokens -> posts (post_id)); diesel::joinable!(postgate_detached -> postgates (postgate_id)); diesel::joinable!(postgate_detached -> posts (detached_post_id)); diesel::joinable!(postgates -> actors (actor_id)); @@ -817,6 +830,7 @@ diesel::allow_tables_to_appear_in_same_query!( post_embed_video_captions, post_facets, post_mentions, + post_search_tokens, postgate_detached, postgates, posts, -- 2.51.2