added recordtrie functions

This commit is contained in:
Soma Nakamura 2025-07-24 04:47:22 +09:00
parent 46a1077d99
commit 78b3c1b633
5 changed files with 618 additions and 1 deletions

52
Cargo.lock generated
View file

@ -11,6 +11,12 @@ dependencies = [
"memchr",
]
[[package]]
name = "base64"
version = "0.21.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567"
[[package]]
name = "bindgen"
version = "0.69.5"
@ -115,6 +121,12 @@ dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
[[package]]
name = "lazy_static"
version = "1.5.0"
@ -159,10 +171,12 @@ checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94"
name = "marisa-rs"
version = "0.1.1"
dependencies = [
"base64",
"bindgen",
"cc",
"libc",
"pkg-config",
"serde_json",
]
[[package]]
@ -275,6 +289,44 @@ dependencies = [
"windows-sys 0.59.0",
]
[[package]]
name = "ryu"
version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
[[package]]
name = "serde"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.141"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30b9eff21ebe718216c6ec64e1d9ac57087aad11efc64e32002bce4a0d4c03d3"
dependencies = [
"itoa",
"memchr",
"ryu",
"serde",
]
[[package]]
name = "shlex"
version = "1.3.0"

View file

@ -16,6 +16,8 @@ crate-type = ["cdylib", "rlib"]
[dependencies]
libc = "0.2"
serde_json = "1.0"
base64 = "0.21"
[build-dependencies]
cc = "1.0"

View file

@ -90,6 +90,47 @@ mmapped_trie.mmap("my_trie.marisa")?;
println!("Trie size: {} bytes", trie.io_size());
```
### RecordTrie Usage
RecordTrie allows storing multiple structured records for each key:
```rust
use marisa_rs::RecordTrie;
// Create a builder
let mut builder = RecordTrie::builder();
// Add structured data (supports duplicate keys)
builder.insert_u32_pair("apple", (1, 100)); // price, quantity
builder.insert_u32_pair("apple", (2, 50)); // different record for same key
builder.insert_u32_pair("banana", (3, 200));
// Add vector data
builder.insert_u32_vec("features", vec![1, 2, 3, 4]);
builder.insert_u32_vec("features", vec![5, 6]); // different length
// Add raw binary data
builder.insert("description", b"Fresh fruit".to_vec());
// Build the RecordTrie
let record_trie = builder.build().unwrap();
// Retrieve all records for a key
let apple_records = record_trie.get_u32_pairs("apple");
println!("Apple records: {:?}", apple_records); // [(1, 100), (2, 50)]
let feature_vecs = record_trie.get_u32_vecs("features");
println!("Feature vectors: {:?}", feature_vecs); // [[1, 2, 3, 4], [5, 6]]
// Prefix search works too
let fruit_keys = record_trie.keys_with_prefix("a");
println!("Keys starting with 'a': {:?}", fruit_keys); // ["apple"]
// Save and load
record_trie.save("trie.marisa", "records.json").unwrap();
let loaded_trie = RecordTrie::load("trie.marisa", "records.json").unwrap();
```
### Lookup Operations
```rust
@ -162,6 +203,24 @@ trie.build(&mut keyset)?;
- `trie.io_size()` - Get the serialized size of the trie in bytes
- `trie.clear()` - Clear the trie, removing all keys (returns `Result<(), &str>`)
### RecordTrie
RecordTrie allows storing structured data associated with keys, similar to Python's marisa-trie RecordTrie:
- `RecordTrie::builder()` - Create a new RecordTrie builder
- `builder.insert(key, data)` - Insert raw binary data for a key
- `builder.insert_u32_pair(key, (a, b))` - Insert a pair of u32 values
- `builder.insert_u32_vec(key, vec![...])` - Insert a vector of u32 values
- `builder.build()` - Build the RecordTrie
- `trie.get(key)` - Get all raw binary records for a key
- `trie.get_u32_pairs(key)` - Get all u32 pairs for a key
- `trie.get_u32_vecs(key)` - Get all u32 vectors for a key
- `trie.contains_key(key)` - Check if a key exists
- `trie.keys_with_prefix(prefix)` - Find all keys with given prefix
- `trie.prefixes_of(query)` - Find all keys that are prefixes of query
- `trie.save(trie_path, records_path)` - Save to files
- `RecordTrie::load(trie_path, records_path)` - Load from files
## Japanese Text Example
```rust

View file

@ -31,6 +31,7 @@
use std::slice;
use std::path::Path;
use std::ffi::CString;
use std::collections::HashMap;
mod bindings {
include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
@ -387,6 +388,264 @@ unsafe impl Send for Keyset {}
unsafe impl Send for Trie {}
unsafe impl Send for Agent {}
/// A RecordTrie for storing structured data associated with keys.
///
/// RecordTrie allows storing multiple structured records for each key,
/// similar to Python's marisa-trie RecordTrie functionality.
///
/// # Example
///
/// ```rust
/// use marisa_rs::RecordTrie;
///
/// let mut record_trie = RecordTrie::new();
/// record_trie.insert("apple", vec![1u32, 2u32]);
/// record_trie.insert("apple", vec![3u32, 4u32]); // Duplicate key
/// record_trie.insert("banana", vec![5u32, 6u32]);
///
/// let trie = record_trie.build().unwrap();
///
/// // Lookup returns all records for a key
/// let apple_records = trie.get("apple");
/// assert_eq!(apple_records.len(), 2); // Two records for "apple"
/// ```
pub struct RecordTrie {
/// Internal trie for key management
trie: Trie,
/// Storage for record data, keyed by trie ID
records: HashMap<usize, Vec<Vec<u8>>>,
}
/// Builder for creating RecordTrie instances
pub struct RecordTrieBuilder {
/// Key-value pairs to be built into the trie
entries: HashMap<String, Vec<Vec<u8>>>,
}
impl RecordTrieBuilder {
/// Creates a new RecordTrieBuilder
pub fn new() -> Self {
RecordTrieBuilder {
entries: HashMap::new(),
}
}
/// Inserts a record for a given key
///
/// Multiple records can be inserted for the same key.
///
/// # Arguments
///
/// * `key` - The string key
/// * `data` - Binary data as Vec<u8>
pub fn insert<K: AsRef<str>>(&mut self, key: K, data: Vec<u8>) {
let key_str = key.as_ref().to_string();
self.entries.entry(key_str).or_insert_with(Vec::new).push(data);
}
/// Inserts a record with automatic serialization for common types
///
/// # Arguments
///
/// * `key` - The string key
/// * `data` - Data that can be serialized to bytes
pub fn insert_u32_pair<K: AsRef<str>>(&mut self, key: K, data: (u32, u32)) {
let mut bytes = Vec::new();
bytes.extend_from_slice(&data.0.to_le_bytes());
bytes.extend_from_slice(&data.1.to_le_bytes());
self.insert(key, bytes);
}
/// Inserts multiple u32 values as a record
pub fn insert_u32_vec<K: AsRef<str>>(&mut self, key: K, data: Vec<u32>) {
let mut bytes = Vec::new();
for value in data {
bytes.extend_from_slice(&value.to_le_bytes());
}
self.insert(key, bytes);
}
/// Builds the RecordTrie
///
/// # Errors
///
/// Returns an error if the trie cannot be built.
pub fn build(self) -> Result<RecordTrie, &'static str> {
// Create keyset with all unique keys
let mut keyset = Keyset::new();
let mut key_order = Vec::new();
for key in self.entries.keys() {
keyset.push(key);
key_order.push(key.clone());
}
// Build the internal trie
let mut trie = Trie::new();
trie.build(&mut keyset)?;
// Create records mapping using trie IDs
let mut records = HashMap::new();
for key in key_order {
if let Some(id) = trie.lookup(&key) {
if let Some(data_list) = self.entries.get(&key) {
records.insert(id, data_list.clone());
}
}
}
Ok(RecordTrie { trie, records })
}
}
impl RecordTrie {
/// Creates a new RecordTrieBuilder for constructing RecordTrie instances
pub fn builder() -> RecordTrieBuilder {
RecordTrieBuilder::new()
}
/// Gets all records associated with a key
///
/// # Arguments
///
/// * `key` - The key to look up
///
/// # Returns
///
/// A vector of all records associated with the key, or an empty vector if the key is not found.
pub fn get<K: AsRef<str>>(&self, key: K) -> Vec<&Vec<u8>> {
if let Some(id) = self.trie.lookup(key.as_ref()) {
if let Some(record_list) = self.records.get(&id) {
return record_list.iter().collect();
}
}
Vec::new()
}
/// Gets records and deserializes them as u32 pairs
pub fn get_u32_pairs<K: AsRef<str>>(&self, key: K) -> Vec<(u32, u32)> {
self.get(key).into_iter()
.filter_map(|bytes| {
if bytes.len() >= 8 {
let first = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
let second = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
Some((first, second))
} else {
None
}
})
.collect()
}
/// Gets records and deserializes them as u32 vectors
pub fn get_u32_vecs<K: AsRef<str>>(&self, key: K) -> Vec<Vec<u32>> {
self.get(key).into_iter()
.map(|bytes| {
bytes.chunks_exact(4)
.map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
.collect()
})
.collect()
}
/// Checks if a key exists in the trie
pub fn contains_key<K: AsRef<str>>(&self, key: K) -> bool {
self.trie.lookup(key.as_ref()).is_some()
}
/// Returns the number of unique keys in the trie
pub fn len(&self) -> usize {
self.trie.size()
}
/// Returns true if the trie is empty
pub fn is_empty(&self) -> bool {
self.trie.is_empty()
}
/// Gets all keys that have the given string as a prefix
pub fn keys_with_prefix<K: AsRef<str>>(&self, prefix: K) -> Vec<String> {
let mut results = Vec::new();
self.trie.predictive_search(prefix.as_ref(), |key, _id| {
results.push(key.to_string());
});
results
}
/// Gets all keys that are prefixes of the given string
pub fn prefixes_of<K: AsRef<str>>(&self, query: K) -> Vec<String> {
let mut results = Vec::new();
self.trie.common_prefix_search(query.as_ref(), |key, _id| {
results.push(key.to_string());
});
results
}
/// Saves the RecordTrie to files
///
/// # Arguments
///
/// * `trie_path` - Path for the trie structure file
/// * `records_path` - Path for the records data file
pub fn save<P: AsRef<Path>>(&self, trie_path: P, records_path: P) -> Result<(), Box<dyn std::error::Error>> {
// Save the internal trie
self.trie.save(&trie_path).map_err(|e| e.to_string())?;
// Convert records to a JSON-serializable format
let mut json_records = std::collections::HashMap::new();
for (id, record_list) in &self.records {
let encoded_records: Vec<String> = record_list.iter()
.map(|bytes| base64::encode_engine(bytes, &base64::engine::general_purpose::STANDARD))
.collect();
json_records.insert(id.to_string(), encoded_records);
}
// Serialize and save records as JSON
let json_data = serde_json::to_string(&json_records)?;
std::fs::write(records_path, json_data)?;
Ok(())
}
/// Loads a RecordTrie from files
///
/// # Arguments
///
/// * `trie_path` - Path to the trie structure file
/// * `records_path` - Path to the records data file
pub fn load<P: AsRef<Path>>(trie_path: P, records_path: P) -> Result<RecordTrie, Box<dyn std::error::Error>> {
// Load the internal trie
let mut trie = Trie::new();
trie.load(trie_path).map_err(|e| e.to_string())?;
// Load and deserialize records from JSON
let json_data = std::fs::read_to_string(records_path)?;
let json_records: std::collections::HashMap<String, Vec<String>> = serde_json::from_str(&json_data)?;
// Convert back to internal format
let mut records = HashMap::new();
for (id_str, encoded_records) in json_records {
let id: usize = id_str.parse()?;
let decoded_records: Result<Vec<Vec<u8>>, _> = encoded_records.iter()
.map(|encoded| base64::decode_engine(encoded, &base64::engine::general_purpose::STANDARD))
.collect();
records.insert(id, decoded_records?);
}
Ok(RecordTrie { trie, records })
}
}
impl Default for RecordTrieBuilder {
fn default() -> Self {
Self::new()
}
}
unsafe impl Send for RecordTrie {}
unsafe impl Send for RecordTrieBuilder {}
#[cfg(test)]
mod tests {
use super::*;
@ -540,4 +799,112 @@ mod tests {
assert!(trie.is_empty());
assert!(trie.lookup("test").is_none());
}
#[test]
fn test_record_trie_basic() {
let mut builder = RecordTrie::builder();
// Add records with duplicate keys
builder.insert_u32_pair("apple", (1, 2));
builder.insert_u32_pair("apple", (3, 4)); // Duplicate key
builder.insert_u32_pair("banana", (5, 6));
let trie = builder.build().expect("Failed to build RecordTrie");
// Test basic functionality
assert_eq!(trie.len(), 2); // Two unique keys
assert!(!trie.is_empty());
assert!(trie.contains_key("apple"));
assert!(trie.contains_key("banana"));
assert!(!trie.contains_key("orange"));
// Test retrieving records
let apple_records = trie.get_u32_pairs("apple");
assert_eq!(apple_records.len(), 2);
assert!(apple_records.contains(&(1, 2)));
assert!(apple_records.contains(&(3, 4)));
let banana_records = trie.get_u32_pairs("banana");
assert_eq!(banana_records.len(), 1);
assert_eq!(banana_records[0], (5, 6));
// Test non-existent key
let orange_records = trie.get_u32_pairs("orange");
assert_eq!(orange_records.len(), 0);
}
#[test]
fn test_record_trie_u32_vectors() {
let mut builder = RecordTrie::builder();
builder.insert_u32_vec("numbers", vec![1, 2, 3, 4]);
builder.insert_u32_vec("numbers", vec![5, 6]); // Different length
builder.insert_u32_vec("single", vec![42]);
let trie = builder.build().expect("Failed to build RecordTrie");
let number_records = trie.get_u32_vecs("numbers");
assert_eq!(number_records.len(), 2);
assert!(number_records.contains(&vec![1, 2, 3, 4]));
assert!(number_records.contains(&vec![5, 6]));
let single_records = trie.get_u32_vecs("single");
assert_eq!(single_records.len(), 1);
assert_eq!(single_records[0], vec![42]);
}
#[test]
fn test_record_trie_raw_bytes() {
let mut builder = RecordTrie::builder();
builder.insert("text", b"hello".to_vec());
builder.insert("text", b"world".to_vec());
builder.insert("binary", vec![0xFF, 0x00, 0xAA, 0x55]);
let trie = builder.build().expect("Failed to build RecordTrie");
let text_records = trie.get("text");
assert_eq!(text_records.len(), 2);
assert!(text_records.iter().any(|&bytes| bytes == b"hello"));
assert!(text_records.iter().any(|&bytes| bytes == b"world"));
let binary_records = trie.get("binary");
assert_eq!(binary_records.len(), 1);
assert_eq!(binary_records[0], &vec![0xFF, 0x00, 0xAA, 0x55]);
}
#[test]
fn test_record_trie_prefix_search() {
let mut builder = RecordTrie::builder();
builder.insert_u32_pair("app", (1, 1));
builder.insert_u32_pair("apple", (2, 2));
builder.insert_u32_pair("application", (3, 3));
builder.insert_u32_pair("banana", (4, 4));
let trie = builder.build().expect("Failed to build RecordTrie");
// Test keys with prefix
let app_keys = trie.keys_with_prefix("app");
assert_eq!(app_keys.len(), 3);
assert!(app_keys.contains(&"app".to_string()));
assert!(app_keys.contains(&"apple".to_string()));
assert!(app_keys.contains(&"application".to_string()));
// Test prefixes of a key
let prefixes = trie.prefixes_of("application");
assert!(prefixes.contains(&"app".to_string()));
assert!(prefixes.contains(&"application".to_string()));
}
#[test]
fn test_record_trie_empty() {
let builder = RecordTrie::builder();
let trie = builder.build().expect("Failed to build empty RecordTrie");
assert_eq!(trie.len(), 0);
assert!(trie.is_empty());
assert!(!trie.contains_key("anything"));
assert_eq!(trie.get("anything").len(), 0);
}
}

View file

@ -1,4 +1,4 @@
use marisa_rs::{Keyset, Trie};
use marisa_rs::{Keyset, Trie, RecordTrie};
fn main() {
println!("Marisa Trie Demo");
@ -49,6 +49,7 @@ fn main() {
});
test_save_load_functionality(&trie);
test_record_trie_functionality();
}
fn test_save_load_functionality(original_trie: &Trie) {
@ -168,3 +169,139 @@ fn test_save_load_functionality(original_trie: &Trie) {
Err(e) => println!("\n⚠ Warning: Could not remove test file '{}': {}", test_file, e),
}
}
fn test_record_trie_functionality() {
println!("\n=== RecordTrie Functionality Test ===");
// Create a RecordTrie with duplicate keys and different data types
let mut builder = RecordTrie::builder();
println!("Building RecordTrie with structured data...");
// Add some records with duplicate keys
builder.insert_u32_pair("apple", (1, 100));
builder.insert_u32_pair("apple", (2, 200)); // Duplicate key
builder.insert_u32_pair("banana", (3, 300));
builder.insert_u32_pair("cherry", (4, 400));
// Add vector data
builder.insert_u32_vec("numbers", vec![1, 2, 3, 4, 5]);
builder.insert_u32_vec("numbers", vec![10, 20]); // Duplicate key, different data
// Add raw binary data
builder.insert("text", b"Hello, World!".to_vec());
builder.insert("text", b"Rust is awesome!".to_vec());
match builder.build() {
Ok(record_trie) => {
println!("✓ RecordTrie built successfully! Keys: {}", record_trie.len());
// Test basic lookup
println!("\n--- Testing Basic Lookup ---");
test_record_lookup(&record_trie);
// Test prefix searches
println!("\n--- Testing Prefix Searches ---");
test_record_prefix_search(&record_trie);
// Test save/load functionality
println!("\n--- Testing RecordTrie Save/Load ---");
test_record_save_load(&record_trie);
}
Err(e) => {
println!("✗ Failed to build RecordTrie: {}", e);
}
}
}
fn test_record_lookup(record_trie: &RecordTrie) {
// Test u32 pairs
let apple_records = record_trie.get_u32_pairs("apple");
println!("'apple' records: {} found", apple_records.len());
for (i, (first, second)) in apple_records.iter().enumerate() {
println!(" Record {}: ({}, {})", i + 1, first, second);
}
let banana_records = record_trie.get_u32_pairs("banana");
println!("'banana' records: {} found", banana_records.len());
for (i, (first, second)) in banana_records.iter().enumerate() {
println!(" Record {}: ({}, {})", i + 1, first, second);
}
// Test u32 vectors
let number_records = record_trie.get_u32_vecs("numbers");
println!("'numbers' records: {} found", number_records.len());
for (i, vec) in number_records.iter().enumerate() {
println!(" Record {}: {:?}", i + 1, vec);
}
// Test raw bytes
let text_records = record_trie.get("text");
println!("'text' records: {} found", text_records.len());
for (i, bytes) in text_records.iter().enumerate() {
if let Ok(text) = std::str::from_utf8(bytes) {
println!(" Record {}: \"{}\"", i + 1, text);
} else {
println!(" Record {}: {:?} (binary)", i + 1, bytes);
}
}
// Test non-existent key
let missing_records = record_trie.get_u32_pairs("missing");
println!("'missing' records: {} found", missing_records.len());
}
fn test_record_prefix_search(record_trie: &RecordTrie) {
// Test prefix search
let keys_with_prefix = record_trie.keys_with_prefix("a");
println!("Keys starting with 'a': {:?}", keys_with_prefix);
// Test finding prefixes
let prefixes = record_trie.prefixes_of("application");
println!("Prefixes of 'application': {:?}", prefixes);
}
fn test_record_save_load(record_trie: &RecordTrie) {
let trie_file = "test_record_trie.marisa";
let records_file = "test_record_data.json";
// Clean up from any previous test
let _ = std::fs::remove_file(trie_file);
let _ = std::fs::remove_file(records_file);
// Test save
match record_trie.save(trie_file, records_file) {
Ok(()) => println!("✓ RecordTrie saved successfully"),
Err(e) => {
println!("✗ Failed to save RecordTrie: {}", e);
return;
}
}
// Test load
match RecordTrie::load(trie_file, records_file) {
Ok(loaded_trie) => {
println!("✓ RecordTrie loaded successfully! Keys: {}", loaded_trie.len());
// Verify the loaded trie works
let apple_records = loaded_trie.get_u32_pairs("apple");
println!(" Verification: 'apple' has {} records after load", apple_records.len());
if apple_records.len() == 2 {
println!("✓ RecordTrie save/load verification passed!");
} else {
println!("⚠ RecordTrie save/load verification failed");
}
}
Err(e) => {
println!("✗ Failed to load RecordTrie: {}", e);
}
}
// Clean up
match (std::fs::remove_file(trie_file), std::fs::remove_file(records_file)) {
(Ok(()), Ok(())) => println!("✓ Cleanup: Removed test files"),
_ => println!("⚠ Warning: Could not remove some test files"),
}
}