1#![deny(missing_docs)]
4
5pub mod catalog;
7pub mod catalog_api;
8mod cluster_state;
9pub mod columnar_api;
10mod dataframe_api;
11pub mod options;
12pub mod owned_session;
13pub mod owned_sql;
14mod sql_api;
15mod txn_manager;
16
17pub use crate::catalog::{CachedTableInfo, Catalog};
18pub use crate::catalog_api::{
19 CatalogInfo, CatalogManifestCatalog, CatalogManifestColumn, CatalogManifestCompression,
20 CatalogManifestDataSourceFormat, CatalogManifestDataType, CatalogManifestDelta,
21 CatalogManifestIndex, CatalogManifestIndexMethod, CatalogManifestNamespace,
22 CatalogManifestRowIdMode, CatalogManifestStorage, CatalogManifestStorageType,
23 CatalogManifestTable, CatalogManifestTableType, CatalogManifestVectorMetric, ColumnDefinition,
24 ColumnInfo, CreateCatalogRequest, CreateNamespaceRequest, CreateTableRequest, IndexInfo,
25 NamespaceInfo, StorageInfo, TableInfo, CATALOG_MANIFEST_DELTA_FORMAT,
26};
27pub use crate::columnar_api::{
28 ColumnarIndexInfo, ColumnarIndexType, ColumnarRowIterator, EmbeddedConfig, StorageMode,
29};
30pub use crate::options::DatabaseOptions;
31pub use crate::owned_session::{EmbeddedOwnedSessionFactory, OwnedEmbeddedTransaction};
32pub use crate::owned_sql::{OwnedSqlRowOutcome, OwnedSqlStreamPlan};
33pub use crate::sql_api::{SqlStreamingResult, StreamingQueryResult, StreamingRows};
34pub use crate::txn_manager::{TransactionInfo, TransactionManager};
35pub use alopex_dataframe::{DataFrame, JoinKeys, JoinType, SortOptions};
36pub use alopex_sql::{DataSourceFormat, TableType};
37pub type SqlResult = alopex_sql::SqlResult;
39pub use alopex_core::kv::{
40 KeyPattern, KeySearchCancellation, KeySearchEntry, KeySearchPage, KeySearchRequest,
41 OwnedReadOptions, OwnedReadSession, OwnedTransactionSession,
42};
43use alopex_core::vector::hnsw::{HnswTransactionState, SearchStats as HnswSearchStats};
44use alopex_core::{
45 columnar::{
46 kvs_bridge::ColumnarKvsBridge, memory::InMemorySegmentStore, segment_v2::SegmentConfigV2,
47 },
48 kv::any::AnyKVTransaction,
49 kv::memory::MemoryKV,
50 kv::{AnyKV, RangeChangeJournalCapability},
51 validate_dimensions, HnswIndex, KVStore, KVTransaction, Key, LargeValueKind, LargeValueMeta,
52 LargeValueReader, LargeValueWriter, StorageFactory, VectorType, DEFAULT_CHUNK_SIZE,
53};
54pub use alopex_core::{HnswConfig, HnswSearchResult, HnswStats, MemoryStats, Metric, TxnMode};
55pub use alopex_sql::executor::QueryRowIterator;
57use alopex_sql::storage::LocalRangeChangeJournal;
58use std::collections::HashMap;
59use std::convert::TryInto;
60use std::fs;
61use std::path::{Path, PathBuf};
62use std::result;
63use std::sync::atomic::{AtomicU64, Ordering};
64use std::sync::{Arc, RwLock};
65
66pub type Result<T> = result::Result<T, Error>;
68
69#[derive(Debug, thiserror::Error)]
71pub enum Error {
72 #[error("core error: {0}")]
74 Core(#[from] alopex_core::Error),
75 #[error("{0}")]
78 ReadAt(#[from] alopex_core::ReadAtError),
79 #[error("{0}")]
81 Sql(#[from] alopex_sql::SqlError),
82 #[error("{0}")]
84 DataFrame(#[from] alopex_dataframe::DataFrameError),
85 #[error("transaction is completed")]
87 TxnCompleted,
88 #[error("カタログが見つかりません: {0}")]
90 CatalogNotFound(String),
91 #[error("カタログは既に存在します: {0}")]
93 CatalogAlreadyExists(String),
94 #[error("カタログが空ではありません: {0}")]
96 CatalogNotEmpty(String),
97 #[error("ネームスペースが見つかりません: {0}.{1}")]
99 NamespaceNotFound(String, String),
100 #[error("ネームスペースは既に存在します: {0}.{1}")]
102 NamespaceAlreadyExists(String, String),
103 #[error("ネームスペースが空ではありません: {0}.{1}")]
105 NamespaceNotEmpty(String, String),
106 #[error("table not found: {0}")]
108 TableNotFound(String),
109 #[error("テーブルは既に存在します: {0}")]
111 TableAlreadyExists(String),
112 #[error("インデックスが見つかりません: {0}")]
114 IndexNotFound(String),
115 #[error("default オブジェクトは削除できません: {0}")]
117 CannotDeleteDefault(String),
118 #[error("managed テーブルにはスキーマが必要です")]
120 SchemaRequired,
121 #[error("external テーブルには storage_root が必要です")]
123 StorageRootRequired,
124 #[error("トランザクションは読み取り専用です")]
126 TxnReadOnly,
127 #[error("invalid transaction id: {0}")]
129 InvalidTransactionId(String),
130 #[error("not in in-memory columnar mode")]
132 NotInMemoryMode,
133 #[error("unsupported data source format: {0}")]
135 UnsupportedDataSourceFormat(String),
136 #[error("catalog lock poisoned")]
138 CatalogLockPoisoned,
139 #[error("cluster state lock poisoned")]
141 ClusterStateLockPoisoned,
142}
143
144impl Error {
145 pub fn sql_error_code(&self) -> Option<&'static str> {
147 match self {
148 Self::Sql(err) => Some(err.code()),
149 _ => None,
150 }
151 }
152}
153
154pub struct Database {
156 pub(crate) store: Arc<AnyKV>,
158 pub(crate) sql_catalog: Arc<RwLock<alopex_sql::catalog::PersistentCatalog<AnyKV>>>,
159 pub(crate) hnsw_cache: RwLock<HashMap<String, Arc<HnswIndex>>>,
160 pub(crate) vector_cache: RwLock<Option<HashMap<Key, CachedVector>>>,
161 pub(crate) table_info_cache: RwLock<HashMap<String, CachedTableInfo>>,
163 pub(crate) table_info_cache_epoch: AtomicU64,
165 pub(crate) columnar_mode: StorageMode,
166 pub(crate) columnar_bridge: ColumnarKvsBridge,
167 pub(crate) columnar_memory: Option<InMemorySegmentStore>,
168 pub(crate) segment_config: SegmentConfigV2,
169 pub(crate) cluster_state: RwLock<cluster_state::EmbeddedClusterState>,
170}
171
172pub(crate) fn disk_data_dir_path(path: &Path) -> std::path::PathBuf {
173 if path.extension().is_some_and(|e| e == "alopex") {
174 path.with_extension("alopex.d")
177 } else {
178 path.to_path_buf()
179 }
180}
181
182#[cfg(not(target_arch = "wasm32"))]
183fn read_file_version_from_storage(path: &Path) -> alopex_core::storage::format::FileVersion {
184 use alopex_core::storage::format::{FileHeader, FileVersion, HEADER_SIZE};
185 use std::io::Read;
186
187 let Some(file_path) = resolve_format_file_path(path) else {
188 return FileVersion::CURRENT;
189 };
190
191 let mut header_bytes = [0u8; HEADER_SIZE];
192 let Ok(mut file) = fs::File::open(file_path) else {
193 return FileVersion::CURRENT;
194 };
195 if file.read_exact(&mut header_bytes).is_err() {
196 return FileVersion::CURRENT;
197 }
198 match FileHeader::from_bytes(&header_bytes) {
199 Ok(header) => header.version,
200 Err(_) => FileVersion::CURRENT,
201 }
202}
203
204#[cfg(target_arch = "wasm32")]
205fn read_file_version_from_storage(_path: &Path) -> alopex_core::storage::format::FileVersion {
206 alopex_core::storage::format::FileVersion::CURRENT
207}
208
209#[cfg(not(target_arch = "wasm32"))]
210fn resolve_format_file_path(path: &Path) -> Option<PathBuf> {
211 if path.is_file() {
212 return Some(path.to_path_buf());
213 }
214
215 if path.is_dir() {
216 if let Some(ext) = path.extension() {
217 if ext == "d" {
218 let candidate = path.with_extension("");
219 if candidate.is_file() {
220 return Some(candidate);
221 }
222 }
223 }
224
225 if let Ok(entries) = fs::read_dir(path) {
226 for entry in entries.flatten() {
227 let entry_path = entry.path();
228 if entry_path.extension().is_some_and(|ext| ext == "alopex") && entry_path.is_file()
229 {
230 return Some(entry_path);
231 }
232 }
233 }
234 }
235
236 None
237}
238
239impl Database {
240 pub fn open(path: &Path) -> Result<Self> {
242 let data_dir = disk_data_dir_path(path);
243 let store = StorageFactory::create(alopex_core::StorageMode::Disk {
244 path: data_dir,
245 config: None,
246 })
247 .map_err(Error::Core)?;
248 let mut db = Self::init(store, StorageMode::Disk, None, SegmentConfigV2::default());
249 db.load_sql_catalog()?;
250 Ok(db)
251 }
252
253 pub fn new() -> Self {
255 let store = AnyKV::Memory(MemoryKV::new());
256 Self::init(
257 store,
258 StorageMode::InMemory,
259 None,
260 SegmentConfigV2::default(),
261 )
262 }
263
264 pub fn open_in_memory() -> Result<Self> {
266 Self::open_in_memory_with_options(DatabaseOptions::in_memory())
267 }
268
269 pub fn open_in_memory_with_options(opts: DatabaseOptions) -> Result<Self> {
271 if !opts.memory_mode() {
272 return Err(Error::Core(alopex_core::Error::InvalidFormat(
273 "memory_mode must be enabled for in-memory open".into(),
274 )));
275 }
276 let store = StorageFactory::create(opts.to_storage_mode(None)).map_err(Error::Core)?;
277 let mut db = Self::init(
278 store,
279 StorageMode::InMemory,
280 opts.memory_limit(),
281 SegmentConfigV2::default(),
282 );
283 db.load_sql_catalog()?;
284 Ok(db)
285 }
286
287 pub fn open_with_uri(uri: &str) -> Result<Self> {
303 if uri.starts_with("s3://") {
305 #[cfg(feature = "s3")]
306 {
307 return Self::open_s3(uri);
308 }
309 #[cfg(not(feature = "s3"))]
310 {
311 return Err(Error::Core(alopex_core::Error::InvalidFormat(
312 "S3 support requires the 's3' feature".into(),
313 )));
314 }
315 }
316
317 let path = if let Some(stripped) = uri.strip_prefix("file://") {
319 stripped
320 } else {
321 uri
322 };
323
324 Self::open(Path::new(path))
325 }
326
327 #[cfg(feature = "s3")]
346 pub fn open_s3(uri: &str) -> Result<Self> {
347 let s3_config = alopex_core::S3Config::from_uri(uri).map_err(Error::Core)?;
348 let store = StorageFactory::create(alopex_core::StorageMode::S3 { config: s3_config })
349 .map_err(Error::Core)?;
350 let mut db = Self::init(store, StorageMode::Disk, None, SegmentConfigV2::default());
351 db.load_sql_catalog()?;
352 Ok(db)
353 }
354
355 pub(crate) fn init(
356 store: AnyKV,
357 columnar_mode: StorageMode,
358 memory_limit: Option<usize>,
359 segment_config: SegmentConfigV2,
360 ) -> Self {
361 let store = Arc::new(store);
362 let sql_catalog = Arc::new(RwLock::new(alopex_sql::catalog::PersistentCatalog::new(
363 store.clone(),
364 )));
365 let columnar_bridge = ColumnarKvsBridge::new(store.clone());
366 let columnar_memory = if matches!(columnar_mode, StorageMode::InMemory) {
367 Some(InMemorySegmentStore::new(memory_limit.map(|v| v as u64)))
368 } else {
369 None
370 };
371
372 Self {
373 store,
374 sql_catalog,
375 hnsw_cache: RwLock::new(HashMap::new()),
376 vector_cache: RwLock::new(None),
377 table_info_cache: RwLock::new(HashMap::new()),
378 table_info_cache_epoch: AtomicU64::new(0),
379 columnar_mode,
380 columnar_bridge,
381 columnar_memory,
382 segment_config,
383 cluster_state: RwLock::new(cluster_state::EmbeddedClusterState::default()),
384 }
385 }
386
387 pub fn cluster_status_snapshot(&self) -> Result<alopex_cluster::ClusterStatusSnapshot> {
389 let state = self
390 .cluster_state
391 .read()
392 .map_err(|_| Error::ClusterStateLockPoisoned)?;
393 Ok(state.status_snapshot(self.table_info_cache_epoch()))
394 }
395
396 pub fn routing_diagnostics(&self) -> Result<alopex_cluster::RoutingDiagnostics> {
398 let state = self
399 .cluster_state
400 .read()
401 .map_err(|_| Error::ClusterStateLockPoisoned)?;
402 Ok(state.routing_diagnostics(self.table_info_cache_epoch()))
403 }
404
405 pub(crate) fn record_routing<C: alopex_sql::Catalog + ?Sized>(
406 &self,
407 catalog: &C,
408 statement: &alopex_sql::Statement,
409 statement_index: usize,
410 ) {
411 let Ok(mut state) = self.cluster_state.write() else {
412 return;
413 };
414 state.record_routing(
415 catalog,
416 statement,
417 statement_index,
418 self.table_info_cache_epoch(),
419 );
420 }
421
422 fn load_sql_catalog(&mut self) -> Result<()> {
423 use alopex_sql::catalog::CatalogError;
424
425 let loaded = match alopex_sql::catalog::PersistentCatalog::load(self.store.clone()) {
426 Ok(catalog) => catalog,
427 Err(CatalogError::Kv(alopex_core::Error::NotFound)) => {
428 alopex_sql::catalog::PersistentCatalog::new(self.store.clone())
429 }
430 Err(other) => return Err(Error::Sql(other.into())),
431 };
432
433 self.sql_catalog = Arc::new(RwLock::new(loaded));
434 Ok(())
435 }
436
437 fn hnsw_cache_get(&self, name: &str) -> Option<Arc<HnswIndex>> {
438 let cache = self.hnsw_cache.read().expect("hnsw cache lock poisoned");
439 cache.get(name).cloned()
440 }
441
442 fn hnsw_cache_insert(&self, name: &str, index: HnswIndex) -> Arc<HnswIndex> {
443 let index = Arc::new(index);
444 let mut cache = self.hnsw_cache.write().expect("hnsw cache lock poisoned");
445 cache.insert(name.to_string(), Arc::clone(&index));
446 index
447 }
448
449 fn hnsw_cache_remove(&self, name: &str) {
450 let mut cache = self.hnsw_cache.write().expect("hnsw cache lock poisoned");
451 cache.remove(name);
452 }
453
454 pub fn table_info_cache_epoch(&self) -> u64 {
456 self.table_info_cache_epoch.load(Ordering::Relaxed)
457 }
458
459 pub fn get_cached_table_info(
461 &self,
462 catalog_name: &str,
463 namespace_name: &str,
464 table_name: &str,
465 ) -> Option<CachedTableInfo> {
466 let cache = self
467 .table_info_cache
468 .read()
469 .expect("table info cache lock poisoned");
470 let key = format!("{}.{}.{}", catalog_name, namespace_name, table_name);
471 cache.get(&key).cloned()
472 }
473
474 pub fn cache_table_info(
476 &self,
477 catalog_name: &str,
478 namespace_name: &str,
479 table_name: &str,
480 info: CachedTableInfo,
481 ) {
482 let mut cache = self
483 .table_info_cache
484 .write()
485 .expect("table info cache lock poisoned");
486 let key = format!("{}.{}.{}", catalog_name, namespace_name, table_name);
487 cache.insert(key, info);
488 }
489
490 pub fn invalidate_table_info_cache(&self) {
492 self.table_info_cache_epoch.fetch_add(1, Ordering::Relaxed);
493 let mut cache = self
494 .table_info_cache
495 .write()
496 .expect("table info cache lock poisoned");
497 cache.clear();
498 }
499
500 pub fn flush(&self) -> Result<()> {
509 self.store.flush().map_err(Error::Core)?;
510 self.store.converge().map_err(Error::Core)
511 }
512
513 pub fn converge(&self) -> Result<()> {
515 self.store.converge().map_err(Error::Core)
516 }
517
518 pub fn close(&self) -> Result<()> {
523 self.store.close().map_err(Error::Core)
524 }
525
526 pub fn container_path(&self) -> Option<PathBuf> {
531 self.store.container_path().map(Path::to_path_buf)
532 }
533
534 pub fn file_format_version(&self) -> alopex_core::storage::format::FileVersion {
536 use alopex_core::storage::format::FileVersion;
537
538 self.store
539 .file_format_storage_dir()
540 .map_or(FileVersion::CURRENT, read_file_version_from_storage)
541 }
542
543 pub fn memory_usage(&self) -> Option<MemoryStats> {
545 match self.store.as_ref() {
546 AnyKV::Memory(kv) => Some(kv.memory_stats()),
547 _ => None,
548 }
549 }
550
551 pub fn persist_to_disk(&self, wal_path: &Path) -> Result<()> {
558 if !matches!(self.store.as_ref(), AnyKV::Memory(_)) {
559 return Err(Error::NotInMemoryMode);
560 }
561 let data_dir = disk_data_dir_path(wal_path);
562 if wal_path.exists() || data_dir.exists() {
563 return Err(Error::Core(alopex_core::Error::PathExists(
564 wal_path.to_path_buf(),
565 )));
566 }
567
568 let tmp_dir = data_dir.with_extension("d.tmp");
572 if tmp_dir.exists() {
573 return Err(Error::Core(alopex_core::Error::PathExists(tmp_dir)));
574 }
575
576 let snapshot = self.snapshot_pairs()?;
577 let is_container_path = wal_path.extension().is_some_and(|e| e == "alopex");
578 let write_result = (|| -> Result<()> {
579 let converge = if is_container_path {
583 alopex_core::lsm::ConvergePolicy::Always {
584 container: wal_path.to_path_buf(),
585 }
586 } else {
587 alopex_core::lsm::ConvergePolicy::Never
588 };
589 let config = alopex_core::lsm::LsmKVConfig {
590 converge,
591 prune_sidecar_on_drop: false,
594 ..alopex_core::lsm::LsmKVConfig::default()
595 };
596
597 let store = StorageFactory::create(alopex_core::StorageMode::Disk {
598 path: tmp_dir.clone(),
599 config: Some(config),
600 })
601 .map_err(Error::Core)?;
602
603 let mut txn = store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
604 for (key, value) in snapshot {
605 txn.put(key, value).map_err(Error::Core)?;
606 }
607 txn.commit_self().map_err(Error::Core)?;
608 store.converge().map_err(Error::Core)?;
609
610 Ok(())
611 })();
612
613 if let Err(e) = write_result {
614 let _ = fs::remove_dir_all(&tmp_dir);
615 let _ = fs::remove_file(wal_path);
616 return Err(e);
617 }
618
619 if let Err(e) = fs::rename(&tmp_dir, &data_dir) {
620 let _ = fs::remove_dir_all(&tmp_dir);
621 let _ = fs::remove_file(wal_path);
622 return Err(Error::Core(e.into()));
623 }
624 Ok(())
625 }
626
627 pub fn clone_to_memory(&self) -> Result<Self> {
629 let snapshot = self.snapshot_pairs()?;
630 let cloned = Database::open_in_memory()?;
631 if snapshot.is_empty() {
632 return Ok(cloned);
633 }
634
635 let mut txn = cloned.begin(TxnMode::ReadWrite)?;
636 for (key, value) in snapshot {
637 txn.put(&key, &value)?;
638 }
639 txn.commit()?;
640 Ok(cloned)
641 }
642
643 pub fn clear(&self) -> Result<()> {
645 let keys: Vec<Key> = self.snapshot_pairs()?.into_iter().map(|(k, _)| k).collect();
646 if keys.is_empty() {
647 return Ok(());
648 }
649 let mut txn = self.begin(TxnMode::ReadWrite)?;
650 for key in keys {
651 txn.delete(&key)?;
652 }
653 txn.commit()
654 }
655
656 pub fn set_memory_limit(&self, bytes: Option<usize>) {
658 if let AnyKV::Memory(kv) = self.store.as_ref() {
659 kv.txn_manager().set_memory_limit(bytes);
660 }
661 }
662
663 pub fn snapshot(&self) -> Vec<(Key, Vec<u8>)> {
665 self.snapshot_pairs().unwrap_or_default()
666 }
667
668 fn snapshot_pairs(&self) -> Result<Vec<(Key, Vec<u8>)>> {
669 let mut txn = self.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
670 let pairs: Vec<(Key, Vec<u8>)> = txn.scan_prefix(b"").map_err(Error::Core)?.collect();
671 txn.commit_self().map_err(Error::Core)?;
672 Ok(pairs)
673 }
674
675 pub fn create_hnsw_index(&self, name: &str, config: HnswConfig) -> Result<()> {
677 let mut txn = self.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
678 let index = HnswIndex::create(name, config).map_err(Error::Core)?;
679 index.save(&mut txn).map_err(Error::Core)?;
680 txn.commit_self().map_err(Error::Core)?;
681 self.hnsw_cache_insert(name, index);
682 Ok(())
683 }
684
685 pub fn drop_hnsw_index(&self, name: &str) -> Result<()> {
687 let mut txn = self.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
688 let index = HnswIndex::load(name, &mut txn).map_err(Error::Core)?;
689 index.drop(&mut txn).map_err(Error::Core)?;
690 txn.commit_self().map_err(Error::Core)?;
691 self.hnsw_cache_remove(name);
692 Ok(())
693 }
694
695 pub fn get_hnsw_stats(&self, name: &str) -> Result<HnswStats> {
697 if let Some(index) = self.hnsw_cache_get(name) {
698 return Ok(index.stats());
699 }
700 let mut txn = self.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
701 let index = HnswIndex::load(name, &mut txn).map_err(Error::Core)?;
702 let stats = index.stats();
703 self.hnsw_cache_insert(name, index);
704 Ok(stats)
705 }
706
707 pub fn compact_hnsw_index(&self, name: &str) -> Result<alopex_core::vector::CompactionResult> {
709 let mut txn = self.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
710 let mut index = HnswIndex::load(name, &mut txn).map_err(Error::Core)?;
711 let result = index.compact().map_err(Error::Core)?;
712 index.save(&mut txn).map_err(Error::Core)?;
713 txn.commit_self().map_err(Error::Core)?;
714 self.hnsw_cache_insert(name, index);
715 Ok(result)
716 }
717
718 pub fn search_hnsw(
720 &self,
721 name: &str,
722 query: &[f32],
723 k: usize,
724 ef_search: Option<usize>,
725 ) -> Result<(Vec<HnswSearchResult>, HnswSearchStats)> {
726 let profile = std::env::var_os("ALOPEX_PROFILE_HNSW").is_some();
727 let total_start = if profile {
728 Some(std::time::Instant::now())
729 } else {
730 None
731 };
732 if let Some(index) = self.hnsw_cache_get(name) {
733 let search_start = if profile {
734 Some(std::time::Instant::now())
735 } else {
736 None
737 };
738 let result = index.search(query, k, ef_search).map_err(Error::Core)?;
739 if let (true, Some(total_start), Some(search_start)) =
740 (profile, total_start, search_start)
741 {
742 let search_time = search_start.elapsed();
743 let total_time = total_start.elapsed();
744 eprintln!(
745 "alopex.hnsw_search cache=hit name={} k={} ef_search={:?} search_ms={:.2} total_ms={:.2}",
746 name,
747 k,
748 ef_search,
749 search_time.as_secs_f64() * 1000.0,
750 total_time.as_secs_f64() * 1000.0
751 );
752 }
753 return Ok(result);
754 }
755
756 let load_start = if profile {
757 Some(std::time::Instant::now())
758 } else {
759 None
760 };
761 let mut txn = self.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
762 let index = HnswIndex::load(name, &mut txn).map_err(Error::Core)?;
763 let load_time = load_start.map(|start| start.elapsed());
764 let index = self.hnsw_cache_insert(name, index);
765 let search_start = if profile {
766 Some(std::time::Instant::now())
767 } else {
768 None
769 };
770 let result = index.search(query, k, ef_search).map_err(Error::Core)?;
771 if let (true, Some(total_start), Some(search_start)) = (profile, total_start, search_start)
772 {
773 let search_time = search_start.elapsed();
774 let total_time = total_start.elapsed();
775 let load_time_ms = load_time
776 .map(|elapsed| elapsed.as_secs_f64() * 1000.0)
777 .unwrap_or(0.0);
778 eprintln!(
779 "alopex.hnsw_search cache=miss name={} k={} ef_search={:?} load_ms={:.2} search_ms={:.2} total_ms={:.2}",
780 name,
781 k,
782 ef_search,
783 load_time_ms,
784 search_time.as_secs_f64() * 1000.0,
785 total_time.as_secs_f64() * 1000.0
786 );
787 }
788 Ok(result)
789 }
790
791 pub fn create_blob_writer(
793 &self,
794 path: &Path,
795 total_len: u64,
796 chunk_size: Option<u32>,
797 ) -> Result<LargeValueWriter> {
798 let meta = LargeValueMeta {
799 kind: LargeValueKind::Blob,
800 total_len,
801 chunk_size: chunk_size.unwrap_or(DEFAULT_CHUNK_SIZE),
802 };
803 LargeValueWriter::create(path, meta).map_err(Error::Core)
804 }
805
806 pub fn create_typed_writer(
808 &self,
809 path: &Path,
810 type_id: u16,
811 total_len: u64,
812 chunk_size: Option<u32>,
813 ) -> Result<LargeValueWriter> {
814 let meta = LargeValueMeta {
815 kind: LargeValueKind::Typed(type_id),
816 total_len,
817 chunk_size: chunk_size.unwrap_or(DEFAULT_CHUNK_SIZE),
818 };
819 LargeValueWriter::create(path, meta).map_err(Error::Core)
820 }
821
822 pub fn open_large_value(&self, path: &Path) -> Result<LargeValueReader> {
824 LargeValueReader::open(path).map_err(Error::Core)
825 }
826
827 pub fn begin(&self, mode: TxnMode) -> Result<Transaction<'_>> {
829 let mut txn = self.store.begin(mode).map_err(Error::Core)?;
830 let journal = if mode == TxnMode::ReadWrite
831 && self.store.range_change_journal_capability()
832 == RangeChangeJournalCapability::Supported
833 {
834 let scope = {
835 let catalog = self.sql_catalog.read().expect("catalog lock poisoned");
836 sql_api::local_journal_scope(&*catalog)
837 };
838 Some(LocalRangeChangeJournal::capture(&mut txn, scope).map_err(Error::Core)?)
839 } else {
840 None
841 };
842 Ok(Transaction {
843 inner: Some(txn),
844 db: self,
845 hnsw_indices: HashMap::new(),
846 overlay: alopex_sql::catalog::CatalogOverlay::new(),
847 vector_cache_updates: HashMap::new(),
848 vector_cache_deletes: Vec::new(),
849 vector_cache_invalidated: false,
850 catalog_modified: false,
851 journal,
852 })
853 }
854}
855
856impl Default for Database {
857 fn default() -> Self {
858 Self::new()
859 }
860}
861
862pub struct Transaction<'a> {
864 inner: Option<AnyKVTransaction<'a>>,
865 db: &'a Database,
866 hnsw_indices: HashMap<String, (HnswIndex, alopex_core::vector::hnsw::HnswTransactionState)>,
867 overlay: alopex_sql::catalog::CatalogOverlay,
868 vector_cache_updates: HashMap<Key, CachedVector>,
869 vector_cache_deletes: Vec<Key>,
870 vector_cache_invalidated: bool,
871 pub(crate) catalog_modified: bool,
873 journal: Option<LocalRangeChangeJournal>,
875}
876
877#[derive(Debug, Clone, PartialEq)]
879pub struct SearchResult {
880 pub key: Key,
882 pub metadata: Vec<u8>,
884 pub score: f32,
886}
887
888const VECTOR_INDEX_KEY: &[u8] = b"__alopex_vector_index";
889
890impl<'a> Transaction<'a> {
891 pub(crate) fn catalog_overlay(&self) -> &alopex_sql::catalog::CatalogOverlay {
892 &self.overlay
893 }
894
895 pub(crate) fn catalog_overlay_mut(&mut self) -> &mut alopex_sql::catalog::CatalogOverlay {
896 &mut self.overlay
897 }
898
899 pub(crate) fn txn_mode(&self) -> Result<TxnMode> {
900 let txn = self.inner.as_ref().ok_or(Error::TxnCompleted)?;
901 Ok(txn.mode())
902 }
903 pub fn get(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>> {
905 self.inner_mut()?.get(&key.to_vec()).map_err(Error::Core)
906 }
907
908 pub fn put(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
910 self.vector_cache_invalidated = true;
911 self.vector_cache_updates.clear();
912 self.vector_cache_deletes.clear();
913 self.inner_mut()?
914 .put(key.to_vec(), value.to_vec())
915 .map_err(Error::Core)
916 }
917
918 pub fn delete(&mut self, key: &[u8]) -> Result<()> {
920 self.vector_cache_deletes.push(key.to_vec());
921 self.inner_mut()?.delete(key.to_vec()).map_err(Error::Core)
922 }
923
924 pub fn scan_prefix(
928 &mut self,
929 prefix: &[u8],
930 ) -> Result<Box<dyn Iterator<Item = (Key, Vec<u8>)> + '_>> {
931 self.inner_mut()?.scan_prefix(prefix).map_err(Error::Core)
932 }
933
934 pub fn search_keys(&mut self, request: &KeySearchRequest) -> Result<KeySearchPage> {
936 self.inner_mut()?.search_keys(request).map_err(Error::Core)
937 }
938
939 pub fn search_keys_with_cancellation(
941 &mut self,
942 request: &KeySearchRequest,
943 cancellation: &KeySearchCancellation,
944 ) -> Result<KeySearchPage> {
945 self.inner_mut()?
946 .search_keys_with_cancellation(request, cancellation)
947 .map_err(Error::Core)
948 }
949
950 pub fn upsert_to_hnsw(
952 &mut self,
953 index_name: &str,
954 key: &[u8],
955 vector: &[f32],
956 metadata: &[u8],
957 ) -> Result<()> {
958 self.ensure_write_txn()?;
959 let (index, state) = self.hnsw_entry_mut(index_name)?;
960 index
961 .upsert_staged(key, vector, metadata, state)
962 .map_err(Error::Core)
963 }
964
965 pub fn delete_from_hnsw(&mut self, index_name: &str, key: &[u8]) -> Result<bool> {
967 self.ensure_write_txn()?;
968 let (index, state) = self.hnsw_entry_mut(index_name)?;
969 index.delete_staged(key, state).map_err(Error::Core)
970 }
971
972 pub fn upsert_vector(
976 &mut self,
977 key: &[u8],
978 metadata: &[u8],
979 vector: &[f32],
980 metric: Metric,
981 ) -> Result<()> {
982 if vector.is_empty() {
983 return Err(Error::Core(alopex_core::Error::InvalidFormat(
984 "vector cannot be empty".into(),
985 )));
986 }
987 let vt = VectorType::new(vector.len(), metric);
988 vt.validate(vector).map_err(Error::Core)?;
989
990 let payload = encode_vector_entry(vt, metadata, vector);
991 let txn = self.inner_mut()?;
992 txn.put(key.to_vec(), payload).map_err(Error::Core)?;
993
994 let mut keys = self.load_vector_index()?;
995 if !keys.iter().any(|k| k == key) {
996 keys.push(key.to_vec());
997 self.persist_vector_index(&keys)?;
998 }
999
1000 let cached = cached_vector_from_entry(metric, metadata.to_vec(), vector.to_vec());
1001 self.vector_cache_updates.insert(key.to_vec(), cached);
1002 self.vector_cache_deletes.retain(|k| k != key);
1003 Ok(())
1004 }
1005
1006 pub fn get_vector(&mut self, key: &[u8], metric: Metric) -> Result<Option<Vec<f32>>> {
1011 let txn = self.inner_mut()?;
1012 let key_vec = key.to_vec();
1013 let Some(raw) = txn.get(&key_vec).map_err(Error::Core)? else {
1014 return Ok(None);
1015 };
1016 let decoded = decode_vector_entry(&raw).map_err(Error::Core)?;
1017 if decoded.metric != metric {
1018 return Err(Error::Core(alopex_core::Error::UnsupportedMetric {
1019 metric: metric.as_str().to_string(),
1020 }));
1021 }
1022 Ok(Some(decoded.vector))
1023 }
1024
1025 pub fn get_vectors(&mut self, keys: &[Key], metric: Metric) -> Result<Vec<Option<Vec<f32>>>> {
1033 let txn = self.inner_mut()?;
1034 let mut out = Vec::with_capacity(keys.len());
1035 for key in keys {
1036 let Some(raw) = txn.get(key).map_err(Error::Core)? else {
1037 out.push(None);
1038 continue;
1039 };
1040 let decoded = decode_vector_entry(&raw).map_err(Error::Core)?;
1041 if decoded.metric != metric {
1042 return Err(Error::Core(alopex_core::Error::UnsupportedMetric {
1043 metric: metric.as_str().to_string(),
1044 }));
1045 }
1046 out.push(Some(decoded.vector));
1047 }
1048 Ok(out)
1049 }
1050
1051 pub fn search_similar(
1056 &mut self,
1057 query_vector: &[f32],
1058 metric: Metric,
1059 top_k: usize,
1060 filter_keys: Option<&[Key]>,
1061 ) -> Result<Vec<SearchResult>> {
1062 if top_k == 0 {
1063 return Ok(Vec::new());
1064 }
1065
1066 let profile = std::env::var_os("ALOPEX_PROFILE_SEARCH_SIMILAR").is_some();
1067 let total_start = if profile {
1068 Some(std::time::Instant::now())
1069 } else {
1070 None
1071 };
1072 let query_norm_sq = query_vector.iter().map(|v| v * v).sum::<f32>();
1073 let query_norm = if matches!(metric, Metric::Cosine) {
1074 query_norm_sq.sqrt()
1075 } else {
1076 0.0
1077 };
1078 let inv_query_norm = if query_norm == 0.0 {
1079 0.0
1080 } else {
1081 1.0 / query_norm
1082 };
1083
1084 if filter_keys.is_none() && self.txn_mode()? == TxnMode::ReadOnly {
1085 let cache = self
1086 .db
1087 .vector_cache
1088 .read()
1089 .expect("vector cache lock poisoned");
1090 if let Some(cache) = cache.as_ref() {
1091 if cache.is_empty() {
1092 return Ok(Vec::new());
1093 }
1094 let keys_len = cache.len();
1095 let mut rows = Vec::with_capacity(keys_len);
1096 let mut score_time = std::time::Duration::ZERO;
1097 for (key, cached) in cache.iter() {
1098 if cached.metric != metric {
1099 return Err(Error::Core(alopex_core::Error::UnsupportedMetric {
1100 metric: metric.as_str().to_string(),
1101 }));
1102 }
1103 validate_dimensions(cached.vector.len(), query_vector.len())
1104 .map_err(Error::Core)?;
1105 let score_start = if profile {
1106 Some(std::time::Instant::now())
1107 } else {
1108 None
1109 };
1110 let dot = dot_product(query_vector, &cached.vector);
1111 let score = match metric {
1112 Metric::Cosine => {
1113 if cached.inv_norm == 0.0 || inv_query_norm == 0.0 {
1114 0.0
1115 } else {
1116 dot * cached.inv_norm * inv_query_norm
1117 }
1118 }
1119 Metric::L2 => {
1120 let dist_sq = query_norm_sq + cached.norm_sq - 2.0 * dot;
1121 -dist_sq.sqrt()
1122 }
1123 Metric::InnerProduct => dot,
1124 };
1125 if let Some(score_start) = score_start {
1126 score_time += score_start.elapsed();
1127 }
1128 rows.push(SearchResult {
1129 key: key.clone(),
1130 metadata: cached.metadata.clone(),
1131 score,
1132 });
1133 }
1134
1135 let rows_total = rows.len();
1136 let sort_start = if profile {
1137 Some(std::time::Instant::now())
1138 } else {
1139 None
1140 };
1141 if rows.len() > top_k {
1142 rows.select_nth_unstable_by(top_k - 1, |a, b| {
1143 b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key))
1144 });
1145 rows.truncate(top_k);
1146 }
1147 rows.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
1148 if let (true, Some(total_start), Some(sort_start)) =
1149 (profile, total_start, sort_start)
1150 {
1151 let sort_time = sort_start.elapsed();
1152 let total_time = total_start.elapsed();
1153 eprintln!(
1154 "alopex.search_similar keys={} results={} top_k={} load_keys_ms={:.2} get_ms={:.2} decode_ms={:.2} score_ms={:.2} sort_ms={:.2} total_ms={:.2}",
1155 keys_len,
1156 rows_total,
1157 top_k,
1158 0.0,
1159 0.0,
1160 0.0,
1161 score_time.as_secs_f64() * 1000.0,
1162 sort_time.as_secs_f64() * 1000.0,
1163 total_time.as_secs_f64() * 1000.0
1164 );
1165 }
1166 return Ok(rows);
1167 }
1168 }
1169
1170 let (keys, load_keys_time) = if profile {
1171 let start = std::time::Instant::now();
1172 let keys = match filter_keys {
1173 Some(keys) => keys.to_vec(),
1174 None => self.load_vector_index()?,
1175 };
1176 (keys, start.elapsed())
1177 } else {
1178 let keys = match filter_keys {
1179 Some(keys) => keys.to_vec(),
1180 None => self.load_vector_index()?,
1181 };
1182 (keys, std::time::Duration::ZERO)
1183 };
1184 if keys.is_empty() {
1185 return Ok(Vec::new());
1186 }
1187
1188 let keys_len = keys.len();
1189 let mut rows = Vec::with_capacity(keys.len());
1190 let txn = self.inner_mut()?;
1191 let mut get_time = std::time::Duration::ZERO;
1192 let mut decode_time = std::time::Duration::ZERO;
1193 let mut score_time = std::time::Duration::ZERO;
1194 if profile {
1195 for key in keys {
1196 let get_start = std::time::Instant::now();
1197 let Some(raw) = txn.get(&key).map_err(Error::Core)? else {
1198 get_time += get_start.elapsed();
1199 continue;
1200 };
1201 get_time += get_start.elapsed();
1202 let decode_start = std::time::Instant::now();
1203 let decoded = decode_vector_entry_view(&raw).map_err(Error::Core)?;
1204 decode_time += decode_start.elapsed();
1205 if decoded.metric != metric {
1206 return Err(Error::Core(alopex_core::Error::UnsupportedMetric {
1207 metric: metric.as_str().to_string(),
1208 }));
1209 }
1210 validate_dimensions(decoded.dim, query_vector.len()).map_err(Error::Core)?;
1211 let score_start = std::time::Instant::now();
1212 let score =
1213 score_from_bytes(metric, query_vector, query_norm, decoded.vector_bytes)?;
1214 score_time += score_start.elapsed();
1215 rows.push(SearchResult {
1216 key,
1217 metadata: decoded.metadata,
1218 score,
1219 });
1220 }
1221 } else {
1222 for key in keys {
1223 let Some(raw) = txn.get(&key).map_err(Error::Core)? else {
1224 continue;
1225 };
1226 let decoded = decode_vector_entry_view(&raw).map_err(Error::Core)?;
1227 if decoded.metric != metric {
1228 return Err(Error::Core(alopex_core::Error::UnsupportedMetric {
1229 metric: metric.as_str().to_string(),
1230 }));
1231 }
1232 validate_dimensions(decoded.dim, query_vector.len()).map_err(Error::Core)?;
1233 let score =
1234 score_from_bytes(metric, query_vector, query_norm, decoded.vector_bytes)?;
1235 rows.push(SearchResult {
1236 key,
1237 metadata: decoded.metadata,
1238 score,
1239 });
1240 }
1241 }
1242
1243 let rows_total = rows.len();
1244 let sort_start = if profile {
1245 Some(std::time::Instant::now())
1246 } else {
1247 None
1248 };
1249 if rows.len() > top_k {
1250 rows.select_nth_unstable_by(top_k - 1, |a, b| {
1251 b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key))
1252 });
1253 rows.truncate(top_k);
1254 }
1255 rows.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key)));
1256 if let (true, Some(total_start), Some(sort_start)) = (profile, total_start, sort_start) {
1257 let sort_time = sort_start.elapsed();
1258 let total_time = total_start.elapsed();
1259 eprintln!(
1260 "alopex.search_similar keys={} results={} top_k={} load_keys_ms={:.2} get_ms={:.2} decode_ms={:.2} score_ms={:.2} sort_ms={:.2} total_ms={:.2}",
1261 keys_len,
1262 rows_total,
1263 top_k,
1264 load_keys_time.as_secs_f64() * 1000.0,
1265 get_time.as_secs_f64() * 1000.0,
1266 decode_time.as_secs_f64() * 1000.0,
1267 score_time.as_secs_f64() * 1000.0,
1268 sort_time.as_secs_f64() * 1000.0,
1269 total_time.as_secs_f64() * 1000.0
1270 );
1271 }
1272 Ok(rows)
1273 }
1274
1275 fn load_vector_index(&mut self) -> Result<Vec<Key>> {
1276 let txn = self.inner_mut()?;
1277 let Some(raw) = txn.get(&VECTOR_INDEX_KEY.to_vec()).map_err(Error::Core)? else {
1278 return Ok(Vec::new());
1279 };
1280 decode_index(&raw).map_err(Error::Core)
1281 }
1282
1283 fn persist_vector_index(&mut self, keys: &[Key]) -> Result<()> {
1284 let txn = self.inner_mut()?;
1285 let encoded = encode_index(keys)?;
1286 txn.put(VECTOR_INDEX_KEY.to_vec(), encoded)
1287 .map_err(Error::Core)
1288 }
1289
1290 pub fn commit(mut self) -> Result<()> {
1292 {
1293 let txn = self.inner.as_mut().ok_or(Error::TxnCompleted)?;
1294 for (index, state) in self.hnsw_indices.values_mut() {
1295 index.commit_staged(txn, state).map_err(Error::Core)?;
1296 }
1297 let mut catalog = self.db.sql_catalog.write().expect("catalog lock poisoned");
1298 catalog
1299 .persist_overlay(txn, &self.overlay)
1300 .map_err(|err| Error::Sql(err.into()))?;
1301 if let Some(journal) = self.journal.take() {
1302 journal.stage(txn).map_err(Error::Core)?;
1303 }
1304
1305 let vector_cache_invalidated = self.vector_cache_invalidated;
1306 let vector_cache_updates = std::mem::take(&mut self.vector_cache_updates);
1307 let vector_cache_deletes = std::mem::take(&mut self.vector_cache_deletes);
1308 if vector_cache_invalidated {
1309 let mut cache = self
1310 .db
1311 .vector_cache
1312 .write()
1313 .expect("vector cache lock poisoned");
1314 *cache = None;
1315 } else if !vector_cache_updates.is_empty() || !vector_cache_deletes.is_empty() {
1316 let needs_rebuild = {
1317 let cache = self
1318 .db
1319 .vector_cache
1320 .read()
1321 .expect("vector cache lock poisoned");
1322 cache.is_none()
1323 };
1324 if needs_rebuild {
1325 let rebuilt = build_vector_cache_from_txn(txn).map_err(Error::Core)?;
1326 let mut cache = self
1327 .db
1328 .vector_cache
1329 .write()
1330 .expect("vector cache lock poisoned");
1331 *cache = Some(rebuilt);
1332 } else {
1333 let mut cache = self
1334 .db
1335 .vector_cache
1336 .write()
1337 .expect("vector cache lock poisoned");
1338 if let Some(cache) = cache.as_mut() {
1339 for key in vector_cache_deletes {
1340 cache.remove(&key);
1341 }
1342 for (key, cached) in vector_cache_updates {
1343 cache.insert(key, cached);
1344 }
1345 }
1346 }
1347 }
1348 }
1349 let txn = self.inner.take().ok_or(Error::TxnCompleted)?;
1350 let hnsw_indices = std::mem::take(&mut self.hnsw_indices);
1351 txn.commit_self().map_err(Error::Core)?;
1352 if !hnsw_indices.is_empty() {
1353 let mut cache = self
1354 .db
1355 .hnsw_cache
1356 .write()
1357 .expect("hnsw cache lock poisoned");
1358 for (name, (index, _state)) in hnsw_indices {
1359 cache.insert(name, Arc::new(index));
1360 }
1361 }
1362
1363 let overlay = std::mem::take(&mut self.overlay);
1365 let catalog_modified = self.catalog_modified;
1366 let mut catalog = self.db.sql_catalog.write().expect("catalog lock poisoned");
1367 catalog.apply_overlay(overlay);
1368 drop(catalog); if catalog_modified {
1371 self.db.invalidate_table_info_cache();
1372 }
1373 Ok(())
1374 }
1375
1376 pub fn rollback_in_place(&mut self) -> Result<()> {
1378 let txn = self.inner.as_mut().ok_or(Error::TxnCompleted)?;
1379 txn.rollback_in_place().map_err(Error::Core)?;
1380 for (index, state) in self.hnsw_indices.values_mut() {
1381 let _ = index.rollback(state);
1382 }
1383 self.hnsw_indices.clear();
1384 self.overlay = alopex_sql::catalog::CatalogOverlay::default();
1385 self.inner = None;
1386 Ok(())
1387 }
1388
1389 pub fn rollback(mut self) -> Result<()> {
1391 if let Some(txn) = self.inner.take() {
1392 for (index, state) in self.hnsw_indices.values_mut() {
1393 let _ = index.rollback(state);
1394 }
1395 self.hnsw_indices.clear();
1396 txn.rollback_self().map_err(Error::Core)
1397 } else {
1398 Err(Error::TxnCompleted)
1399 }
1400 }
1401
1402 fn inner_mut(&mut self) -> Result<&mut AnyKVTransaction<'a>> {
1403 self.inner.as_mut().ok_or(Error::TxnCompleted)
1404 }
1405
1406 fn hnsw_entry_mut(&mut self, name: &str) -> Result<&mut (HnswIndex, HnswTransactionState)> {
1407 if !self.hnsw_indices.contains_key(name) {
1408 let index = {
1409 let txn = self.inner_mut()?;
1410 HnswIndex::load(name, txn).map_err(Error::Core)?
1411 };
1412 self.hnsw_indices
1413 .insert(name.to_string(), (index, HnswTransactionState::default()));
1414 }
1415 Ok(self.hnsw_indices.get_mut(name).unwrap())
1416 }
1417
1418 fn ensure_write_txn(&self) -> Result<()> {
1419 let txn = self.inner.as_ref().ok_or(Error::TxnCompleted)?;
1420 if txn.mode() != TxnMode::ReadWrite {
1421 return Err(Error::Core(alopex_core::Error::TxnReadOnly));
1422 }
1423 Ok(())
1424 }
1425}
1426
1427impl OwnedEmbeddedTransaction {
1428 pub fn upsert_vector(
1430 &mut self,
1431 key: &[u8],
1432 metadata: &[u8],
1433 vector: &[f32],
1434 metric: Metric,
1435 ) -> Result<()> {
1436 if vector.is_empty() {
1437 return Err(Error::Core(alopex_core::Error::InvalidFormat(
1438 "vector cannot be empty".into(),
1439 )));
1440 }
1441 let vector_type = VectorType::new(vector.len(), metric);
1442 vector_type.validate(vector).map_err(Error::Core)?;
1443 let key = key.to_vec();
1444 let payload = encode_vector_entry(vector_type, metadata, vector);
1445
1446 self.session
1447 .with_transaction(|transaction| {
1448 transaction.put(key.clone(), payload)?;
1449 let mut keys = match transaction.get(&VECTOR_INDEX_KEY.to_vec())? {
1450 Some(raw) => decode_index(&raw)?,
1451 None => Vec::new(),
1452 };
1453 if !keys.iter().any(|entry| entry == &key) {
1454 keys.push(key.clone());
1455 transaction.put(VECTOR_INDEX_KEY.to_vec(), encode_index(&keys)?)?;
1456 }
1457 Ok(())
1458 })
1459 .map_err(Error::Core)?;
1460 self.vector_cache_invalidated = true;
1461 Ok(())
1462 }
1463
1464 pub fn get_vector(&mut self, key: &[u8], metric: Metric) -> Result<Option<Vec<f32>>> {
1466 self.session
1467 .with_transaction(|transaction| {
1468 let Some(raw) = transaction.get(&key.to_vec())? else {
1469 return Ok(None);
1470 };
1471 let decoded = decode_vector_entry(&raw)?;
1472 if decoded.metric != metric {
1473 return Err(alopex_core::Error::UnsupportedMetric {
1474 metric: metric.as_str().to_string(),
1475 });
1476 }
1477 Ok(Some(decoded.vector))
1478 })
1479 .map_err(Error::Core)
1480 }
1481
1482 pub fn get_vectors(&mut self, keys: &[Key], metric: Metric) -> Result<Vec<Option<Vec<f32>>>> {
1484 self.session
1485 .with_transaction(|transaction| {
1486 let mut output = Vec::with_capacity(keys.len());
1487 for key in keys {
1488 let Some(raw) = transaction.get(key)? else {
1489 output.push(None);
1490 continue;
1491 };
1492 let decoded = decode_vector_entry(&raw)?;
1493 if decoded.metric != metric {
1494 return Err(alopex_core::Error::UnsupportedMetric {
1495 metric: metric.as_str().to_string(),
1496 });
1497 }
1498 output.push(Some(decoded.vector));
1499 }
1500 Ok(output)
1501 })
1502 .map_err(Error::Core)
1503 }
1504
1505 pub fn search_similar(
1507 &mut self,
1508 query_vector: &[f32],
1509 metric: Metric,
1510 top_k: usize,
1511 filter_keys: Option<&[Key]>,
1512 ) -> Result<Vec<SearchResult>> {
1513 if top_k == 0 {
1514 return Ok(Vec::new());
1515 }
1516 let query_norm_sq = query_vector.iter().map(|value| value * value).sum::<f32>();
1517 let query_norm = if metric == Metric::Cosine {
1518 query_norm_sq.sqrt()
1519 } else {
1520 0.0
1521 };
1522 let keys = match filter_keys {
1523 Some(keys) => keys.to_vec(),
1524 None => self
1525 .session
1526 .with_transaction(|transaction| {
1527 match transaction.get(&VECTOR_INDEX_KEY.to_vec())? {
1528 Some(raw) => decode_index(&raw),
1529 None => Ok(Vec::new()),
1530 }
1531 })
1532 .map_err(Error::Core)?,
1533 };
1534 let mut rows = self
1535 .session
1536 .with_transaction(|transaction| {
1537 let mut rows = Vec::with_capacity(keys.len());
1538 for key in &keys {
1539 let Some(raw) = transaction.get(key)? else {
1540 continue;
1541 };
1542 let decoded = decode_vector_entry_view(&raw)?;
1543 if decoded.metric != metric {
1544 return Err(alopex_core::Error::UnsupportedMetric {
1545 metric: metric.as_str().to_string(),
1546 });
1547 }
1548 validate_dimensions(decoded.dim, query_vector.len())?;
1549 let score =
1550 score_from_bytes(metric, query_vector, query_norm, decoded.vector_bytes)?;
1551 rows.push(SearchResult {
1552 key: key.clone(),
1553 metadata: decoded.metadata,
1554 score,
1555 });
1556 }
1557 Ok(rows)
1558 })
1559 .map_err(Error::Core)?;
1560 if rows.len() > top_k {
1561 rows.select_nth_unstable_by(top_k - 1, |left, right| {
1562 right
1563 .score
1564 .total_cmp(&left.score)
1565 .then_with(|| left.key.cmp(&right.key))
1566 });
1567 rows.truncate(top_k);
1568 }
1569 rows.sort_by(|left, right| {
1570 right
1571 .score
1572 .total_cmp(&left.score)
1573 .then_with(|| left.key.cmp(&right.key))
1574 });
1575 Ok(rows)
1576 }
1577
1578 pub fn upsert_to_hnsw(
1580 &mut self,
1581 index_name: &str,
1582 key: &[u8],
1583 vector: &[f32],
1584 metadata: &[u8],
1585 ) -> Result<()> {
1586 self.ensure_owned_write_transaction()?;
1587 let (index, state) = self.hnsw_entry_mut(index_name)?;
1588 index
1589 .upsert_staged(key, vector, metadata, state)
1590 .map_err(Error::Core)
1591 }
1592
1593 pub fn delete_from_hnsw(&mut self, index_name: &str, key: &[u8]) -> Result<bool> {
1595 self.ensure_owned_write_transaction()?;
1596 let (index, state) = self.hnsw_entry_mut(index_name)?;
1597 index.delete_staged(key, state).map_err(Error::Core)
1598 }
1599
1600 fn ensure_owned_write_transaction(&self) -> Result<()> {
1601 let mode = self
1602 .session
1603 .with_transaction(|transaction| Ok(transaction.mode()))
1604 .map_err(Error::Core)?;
1605 if mode == TxnMode::ReadOnly {
1606 return Err(Error::Core(alopex_core::Error::TxnReadOnly));
1607 }
1608 Ok(())
1609 }
1610
1611 fn hnsw_entry_mut(
1612 &mut self,
1613 index_name: &str,
1614 ) -> Result<&mut (HnswIndex, HnswTransactionState)> {
1615 if !self.hnsw_indices.contains_key(index_name) {
1616 let index = self
1617 .session
1618 .with_transaction(|transaction| {
1619 let mut transaction =
1620 alopex_core::kv::OwnedKVTransactionAdapter::new(transaction);
1621 HnswIndex::load(index_name, &mut transaction)
1622 })
1623 .map_err(Error::Core)?;
1624 self.hnsw_indices.insert(
1625 index_name.to_string(),
1626 (index, HnswTransactionState::default()),
1627 );
1628 }
1629 Ok(self
1630 .hnsw_indices
1631 .get_mut(index_name)
1632 .expect("HNSW entry inserted above"))
1633 }
1634}
1635
1636impl<'a> Drop for Transaction<'a> {
1637 fn drop(&mut self) {
1638 if let Some(txn) = self.inner.take() {
1639 for (index, state) in self.hnsw_indices.values_mut() {
1640 let _ = index.rollback(state);
1641 }
1642 self.hnsw_indices.clear();
1643 let _ = txn.rollback_self();
1644 }
1645 }
1646}
1647
1648fn metric_to_byte(metric: Metric) -> u8 {
1649 match metric {
1650 Metric::Cosine => 0,
1651 Metric::L2 => 1,
1652 Metric::InnerProduct => 2,
1653 }
1654}
1655
1656fn byte_to_metric(byte: u8) -> result::Result<Metric, alopex_core::Error> {
1657 match byte {
1658 0 => Ok(Metric::Cosine),
1659 1 => Ok(Metric::L2),
1660 2 => Ok(Metric::InnerProduct),
1661 other => Err(alopex_core::Error::UnsupportedMetric {
1662 metric: format!("unknown({other})"),
1663 }),
1664 }
1665}
1666
1667fn encode_vector_entry(vector_type: VectorType, metadata: &[u8], vector: &[f32]) -> Vec<u8> {
1668 let dim = vector_type.dim() as u32;
1669 let meta_len = metadata.len() as u32;
1670 let mut buf = Vec::with_capacity(1 + 4 + 4 + metadata.len() + std::mem::size_of_val(vector));
1671 buf.push(metric_to_byte(vector_type.metric()));
1672 buf.extend_from_slice(&dim.to_le_bytes());
1673 buf.extend_from_slice(&meta_len.to_le_bytes());
1674 buf.extend_from_slice(metadata);
1675 for v in vector {
1676 buf.extend_from_slice(&v.to_le_bytes());
1677 }
1678 buf
1679}
1680
1681struct DecodedEntry {
1682 metric: Metric,
1683 vector: Vec<f32>,
1684}
1685
1686#[derive(Clone)]
1687struct CachedVector {
1688 metric: Metric,
1689 metadata: Vec<u8>,
1690 vector: Vec<f32>,
1691 norm_sq: f32,
1692 inv_norm: f32,
1693}
1694
1695struct VectorEntryView<'a> {
1696 metric: Metric,
1697 dim: usize,
1698 metadata: Vec<u8>,
1699 vector_bytes: &'a [u8],
1700}
1701
1702fn decode_vector_entry(bytes: &[u8]) -> result::Result<DecodedEntry, alopex_core::Error> {
1703 if bytes.len() < 9 {
1704 return Err(alopex_core::Error::InvalidFormat(
1705 "vector entry too short".into(),
1706 ));
1707 }
1708 let metric = byte_to_metric(bytes[0])?;
1709 let dim = u32::from_le_bytes(bytes[1..5].try_into().unwrap()) as usize;
1710 let meta_len = u32::from_le_bytes(bytes[5..9].try_into().unwrap()) as usize;
1711
1712 let header = 9;
1713 let expected_len = header + meta_len + dim * std::mem::size_of::<f32>();
1714 if bytes.len() < expected_len {
1715 return Err(alopex_core::Error::InvalidFormat(
1716 "vector entry truncated".into(),
1717 ));
1718 }
1719
1720 let mut vector = Vec::with_capacity(dim);
1721 let vec_bytes = &bytes[header + meta_len..expected_len];
1722 for chunk in vec_bytes.as_chunks::<4>().0 {
1723 vector.push(f32::from_le_bytes(*chunk));
1724 }
1725
1726 Ok(DecodedEntry { metric, vector })
1727}
1728
1729fn decode_vector_entry_view(
1730 bytes: &[u8],
1731) -> result::Result<VectorEntryView<'_>, alopex_core::Error> {
1732 if bytes.len() < 9 {
1733 return Err(alopex_core::Error::InvalidFormat(
1734 "vector entry too short".into(),
1735 ));
1736 }
1737 let metric = byte_to_metric(bytes[0])?;
1738 let dim = u32::from_le_bytes(bytes[1..5].try_into().unwrap()) as usize;
1739 let meta_len = u32::from_le_bytes(bytes[5..9].try_into().unwrap()) as usize;
1740
1741 let header = 9;
1742 let expected_len = header + meta_len + dim * std::mem::size_of::<f32>();
1743 if bytes.len() < expected_len {
1744 return Err(alopex_core::Error::InvalidFormat(
1745 "vector entry truncated".into(),
1746 ));
1747 }
1748
1749 let metadata = bytes[header..header + meta_len].to_vec();
1750 let vector_bytes = &bytes[header + meta_len..expected_len];
1751
1752 Ok(VectorEntryView {
1753 metric,
1754 dim,
1755 metadata,
1756 vector_bytes,
1757 })
1758}
1759
1760fn vector_bytes_to_vec(bytes: &[u8]) -> Vec<f32> {
1761 let mut vector = Vec::with_capacity(bytes.len() / 4);
1762 for chunk in bytes.as_chunks::<4>().0 {
1763 vector.push(f32::from_le_bytes(*chunk));
1764 }
1765 vector
1766}
1767
1768fn cached_vector_from_entry(metric: Metric, metadata: Vec<u8>, vector: Vec<f32>) -> CachedVector {
1769 let norm_sq = vector.iter().map(|v| v * v).sum::<f32>();
1770 let inv_norm = if norm_sq == 0.0 {
1771 0.0
1772 } else {
1773 1.0 / norm_sq.sqrt()
1774 };
1775 CachedVector {
1776 metric,
1777 metadata,
1778 vector,
1779 norm_sq,
1780 inv_norm,
1781 }
1782}
1783
1784fn build_vector_cache_from_txn<'a>(
1785 txn: &mut AnyKVTransaction<'a>,
1786) -> result::Result<HashMap<Key, CachedVector>, alopex_core::Error> {
1787 let Some(raw) = txn.get(&VECTOR_INDEX_KEY.to_vec())? else {
1788 return Ok(HashMap::new());
1789 };
1790 let keys = decode_index(&raw)?;
1791 let mut cache = HashMap::with_capacity(keys.len());
1792 for key in keys {
1793 let Some(raw) = txn.get(&key)? else {
1794 continue;
1795 };
1796 let decoded = decode_vector_entry_view(&raw)?;
1797 let vector = vector_bytes_to_vec(decoded.vector_bytes);
1798 let cached = cached_vector_from_entry(decoded.metric, decoded.metadata, vector);
1799 cache.insert(key, cached);
1800 }
1801 Ok(cache)
1802}
1803
1804fn dot_product(query: &[f32], item: &[f32]) -> f32 {
1805 #[cfg(target_arch = "x86_64")]
1806 {
1807 if std::is_x86_feature_detected!("avx") {
1808 unsafe {
1810 return dot_product_avx(query, item);
1811 }
1812 }
1813 }
1814 dot_product_scalar(query, item)
1815}
1816
1817fn dot_product_scalar(query: &[f32], item: &[f32]) -> f32 {
1818 query.iter().zip(item.iter()).map(|(q, v)| q * v).sum()
1819}
1820
1821#[cfg(target_arch = "x86_64")]
1822#[target_feature(enable = "avx")]
1823unsafe fn dot_product_avx(query: &[f32], item: &[f32]) -> f32 {
1824 use std::arch::x86_64::*;
1825
1826 let len = query.len();
1827 let mut i = 0;
1828 let mut acc = _mm256_setzero_ps();
1829 let q_ptr = query.as_ptr();
1830 let v_ptr = item.as_ptr();
1831 while i + 8 <= len {
1832 let q = _mm256_loadu_ps(q_ptr.add(i));
1833 let v = _mm256_loadu_ps(v_ptr.add(i));
1834 acc = _mm256_add_ps(acc, _mm256_mul_ps(q, v));
1835 i += 8;
1836 }
1837
1838 let mut tmp = [0f32; 8];
1839 _mm256_storeu_ps(tmp.as_mut_ptr(), acc);
1840 let mut sum = tmp.iter().sum::<f32>();
1841 while i < len {
1842 sum += *q_ptr.add(i) * *v_ptr.add(i);
1843 i += 1;
1844 }
1845 sum
1846}
1847
1848fn score_from_slice(metric: Metric, query: &[f32], query_norm: f32, item: &[f32]) -> f32 {
1849 match metric {
1850 Metric::Cosine => {
1851 if query_norm == 0.0 {
1852 return 0.0;
1853 }
1854 let mut dot = 0.0;
1855 let mut item_norm_sq = 0.0;
1856 for (q, v) in query.iter().zip(item.iter()) {
1857 dot += q * v;
1858 item_norm_sq += v * v;
1859 }
1860 let item_norm = item_norm_sq.sqrt();
1861 if item_norm == 0.0 {
1862 0.0
1863 } else {
1864 dot / (query_norm * item_norm)
1865 }
1866 }
1867 Metric::L2 => {
1868 let mut dist_sq = 0.0;
1869 for (q, v) in query.iter().zip(item.iter()) {
1870 let d = q - v;
1871 dist_sq += d * d;
1872 }
1873 -dist_sq.sqrt()
1874 }
1875 Metric::InnerProduct => query.iter().zip(item.iter()).map(|(q, v)| q * v).sum(),
1876 }
1877}
1878
1879fn score_from_bytes(
1880 metric: Metric,
1881 query: &[f32],
1882 query_norm: f32,
1883 vector_bytes: &[u8],
1884) -> result::Result<f32, alopex_core::Error> {
1885 let len = vector_bytes.len() / 4;
1886 #[cfg(target_endian = "little")]
1887 {
1888 let ptr = vector_bytes.as_ptr();
1889 if (ptr as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1890 let items = unsafe { std::slice::from_raw_parts(ptr as *const f32, len) };
1891 return Ok(score_from_slice(metric, query, query_norm, items));
1892 }
1893 }
1894
1895 let mut iter = vector_bytes.as_chunks::<4>().0.iter();
1897 let score = match metric {
1898 Metric::Cosine => {
1899 if query_norm == 0.0 {
1900 0.0
1901 } else {
1902 let mut dot = 0.0;
1903 let mut item_norm_sq = 0.0;
1904 for (q, chunk) in query.iter().zip(&mut iter) {
1905 let v = f32::from_le_bytes(*chunk);
1906 dot += q * v;
1907 item_norm_sq += v * v;
1908 }
1909 let item_norm = item_norm_sq.sqrt();
1910 if item_norm == 0.0 {
1911 0.0
1912 } else {
1913 dot / (query_norm * item_norm)
1914 }
1915 }
1916 }
1917 Metric::L2 => {
1918 let mut dist_sq = 0.0;
1919 for (q, chunk) in query.iter().zip(&mut iter) {
1920 let v = f32::from_le_bytes(*chunk);
1921 let d = q - v;
1922 dist_sq += d * d;
1923 }
1924 -dist_sq.sqrt()
1925 }
1926 Metric::InnerProduct => query
1927 .iter()
1928 .zip(&mut iter)
1929 .map(|(q, chunk)| q * f32::from_le_bytes(*chunk))
1930 .sum(),
1931 };
1932 Ok(score)
1933}
1934
1935fn encode_index(keys: &[Key]) -> result::Result<Vec<u8>, alopex_core::Error> {
1936 let mut buf = Vec::new();
1937 let count = keys.len() as u32;
1938 buf.extend_from_slice(&count.to_le_bytes());
1939 for key in keys {
1940 let len: u32 = key
1941 .len()
1942 .try_into()
1943 .map_err(|_| alopex_core::Error::InvalidFormat("key too long".into()))?;
1944 buf.extend_from_slice(&len.to_le_bytes());
1945 buf.extend_from_slice(key);
1946 }
1947 Ok(buf)
1948}
1949
1950fn decode_index(bytes: &[u8]) -> result::Result<Vec<Key>, alopex_core::Error> {
1951 if bytes.len() < 4 {
1952 return Err(alopex_core::Error::InvalidFormat("index too short".into()));
1953 }
1954 let count = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
1955 let mut pos = 4;
1956 let mut keys = Vec::with_capacity(count);
1957 for _ in 0..count {
1958 if pos + 4 > bytes.len() {
1959 return Err(alopex_core::Error::InvalidFormat("index truncated".into()));
1960 }
1961 let len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
1962 pos += 4;
1963 if pos + len > bytes.len() {
1964 return Err(alopex_core::Error::InvalidFormat(
1965 "index key truncated".into(),
1966 ));
1967 }
1968 keys.push(bytes[pos..pos + len].to_vec());
1969 pos += len;
1970 }
1971 Ok(keys)
1972}
1973
1974#[cfg(test)]
1975mod tests {
1976 use super::*;
1977 use std::sync::mpsc;
1978 use std::thread;
1979 use tempfile::tempdir;
1980
1981 #[test]
1982 fn test_open_and_crud() {
1983 let dir = tempdir().unwrap();
1984 let path = dir.path().join("test.db");
1985 let db = Database::open(&path).unwrap();
1986
1987 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
1988 txn.put(b"key1", b"value1").unwrap();
1989 txn.commit().unwrap();
1990
1991 let mut txn2 = db.begin(TxnMode::ReadOnly).unwrap();
1992 let val = txn2.get(b"key1").unwrap();
1993 assert_eq!(val, Some(b"value1".to_vec()));
1994 }
1995
1996 #[test]
1997 fn test_not_found() {
1998 let db = Database::new();
1999 let mut txn = db.begin(TxnMode::ReadOnly).unwrap();
2000 let val = txn.get(b"non-existent-key").unwrap();
2001 assert!(val.is_none());
2002 }
2003
2004 #[cfg(not(target_arch = "wasm32"))]
2005 #[test]
2006 fn test_file_format_version_reads_alopex_header() {
2007 use alopex_core::storage::format::{AlopexFileWriter, FileFlags, FileVersion};
2008
2009 let dir = tempdir().unwrap();
2010 let path = dir.path().join("format-test.alopex");
2011 let expected = FileVersion::new(0, 0, 1);
2012
2013 let writer = AlopexFileWriter::new(path.clone(), expected, FileFlags(0)).unwrap();
2014 writer.finalize().unwrap();
2015
2016 let db = Database::open(&path).unwrap();
2017 assert_eq!(db.file_format_version(), expected);
2018 }
2019
2020 #[test]
2021 fn test_crash_recovery_replays_wal() {
2022 let dir = tempdir().unwrap();
2023 let path = dir.path().join("replay.db");
2024
2025 {
2026 let db = Database::open(&path).unwrap();
2027 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2028 txn.put(b"k1", b"v1").unwrap();
2029 txn.commit().unwrap();
2030
2031 let mut uncommitted = db.begin(TxnMode::ReadWrite).unwrap();
2032 uncommitted.put(b"k2", b"v2").unwrap();
2033 }
2035
2036 let db = Database::open(&path).unwrap();
2037 let mut txn = db.begin(TxnMode::ReadOnly).unwrap();
2038 assert_eq!(txn.get(b"k1").unwrap(), Some(b"v1".to_vec()));
2039 assert_eq!(txn.get(b"k2").unwrap(), None);
2040 }
2041
2042 #[test]
2043 fn test_txn_closed() {
2044 let db = Database::new();
2045 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2046 txn.put(b"k1", b"v1").unwrap();
2047 txn.commit().unwrap();
2048 }
2054
2055 #[test]
2056 fn test_concurrency_conflict() {
2057 let db = std::sync::Arc::new(Database::new());
2058 let mut t0 = db.begin(TxnMode::ReadWrite).unwrap();
2059 t0.put(b"k1", b"v0").unwrap();
2060 t0.commit().unwrap();
2061
2062 let (tx1, rx1) = mpsc::channel();
2063 let (tx2, rx2) = mpsc::channel();
2064
2065 let db1 = db.clone();
2066 let t1 = thread::spawn(move || {
2067 let mut txn1 = db1.begin(TxnMode::ReadWrite).unwrap();
2068 let val = txn1.get(b"k1").unwrap();
2069 assert_eq!(val.unwrap(), b"v0");
2070 tx1.send(()).unwrap();
2071 rx2.recv().unwrap();
2072 txn1.put(b"k1", b"v1").unwrap();
2073 let result = txn1.commit();
2074 assert!(matches!(
2075 result,
2076 Err(Error::Core(alopex_core::Error::TxnConflict))
2077 ));
2078 });
2079
2080 let db2 = db.clone();
2081 let t2 = thread::spawn(move || {
2082 rx1.recv().unwrap();
2083 let mut txn2 = db2.begin(TxnMode::ReadWrite).unwrap();
2084 txn2.put(b"k1", b"v2").unwrap();
2085 assert!(txn2.commit().is_ok());
2086 tx2.send(()).unwrap();
2087 });
2088
2089 t1.join().unwrap();
2090 t2.join().unwrap();
2091
2092 let mut txn3 = db.begin(TxnMode::ReadOnly).unwrap();
2093 let val = txn3.get(b"k1").unwrap();
2094 assert_eq!(val.unwrap(), b"v2");
2095 }
2096
2097 #[test]
2098 fn test_flush_and_reopen_via_embedded_api() {
2099 let dir = tempdir().unwrap();
2100 let path = dir.path().join("persist.db");
2101 {
2102 let db = Database::open(&path).unwrap();
2103 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2104 txn.put(b"k1", b"v1").unwrap();
2105 txn.commit().unwrap();
2106 db.flush().unwrap();
2107 }
2108
2109 let db = Database::open(&path).unwrap();
2110 let mut txn = db.begin(TxnMode::ReadOnly).unwrap();
2111 assert_eq!(txn.get(b"k1").unwrap(), Some(b"v1".to_vec()));
2112 }
2113
2114 #[test]
2115 fn test_large_value_blob_roundtrip() {
2116 let dir = tempdir().unwrap();
2117 let path = dir.path().join("blob.lv");
2118 let payload = b"hello large value";
2119
2120 {
2121 let db = Database::new();
2122 let mut writer = db
2123 .create_blob_writer(&path, payload.len() as u64, Some(16))
2124 .unwrap();
2125 writer.write_chunk(&payload[..5]).unwrap();
2126 writer.write_chunk(&payload[5..]).unwrap();
2127 writer.finish().unwrap();
2128 }
2129
2130 let db = Database::new();
2131 let mut reader = db.open_large_value(&path).unwrap();
2132 let mut buf = Vec::new();
2133 while let Some((_info, chunk)) = reader.next_chunk().unwrap() {
2134 buf.extend_from_slice(&chunk);
2135 }
2136 assert_eq!(buf, payload);
2137 }
2138
2139 #[test]
2140 fn upsert_and_search_same_txn() {
2141 let db = Database::new();
2142 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2143 txn.upsert_vector(b"k1", b"meta1", &[1.0, 0.0], Metric::Cosine)
2144 .unwrap();
2145
2146 let results = txn
2147 .search_similar(&[1.0, 0.0], Metric::Cosine, 1, None)
2148 .unwrap();
2149 assert_eq!(results.len(), 1);
2150 assert_eq!(results[0].key, b"k1");
2151 assert_eq!(results[0].metadata, b"meta1");
2152 txn.commit().unwrap();
2153 }
2154
2155 #[test]
2156 fn upsert_and_search_across_txn() {
2157 let db = Database::new();
2158 {
2159 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2160 txn.upsert_vector(b"k1", b"meta1", &[1.0, 1.0], Metric::Cosine)
2161 .unwrap();
2162 txn.commit().unwrap();
2163 }
2164
2165 let mut ro = db.begin(TxnMode::ReadOnly).unwrap();
2166 let results = ro
2167 .search_similar(&[1.0, 1.0], Metric::Cosine, 1, None)
2168 .unwrap();
2169 assert_eq!(results.len(), 1);
2170 assert_eq!(results[0].key, b"k1");
2171 }
2172
2173 #[test]
2174 fn read_only_upsert_rejected() {
2175 let db = Database::new();
2176 let mut ro = db.begin(TxnMode::ReadOnly).unwrap();
2177 let err = ro
2178 .upsert_vector(b"k1", b"m", &[1.0, 0.0], Metric::Cosine)
2179 .unwrap_err();
2180 assert!(matches!(err, Error::Core(alopex_core::Error::TxnReadOnly)));
2181 }
2182
2183 #[test]
2184 fn dimension_mismatch_on_search() {
2185 let db = Database::new();
2186 {
2187 let mut txn = db.begin(TxnMode::ReadWrite).unwrap();
2188 txn.upsert_vector(b"k1", b"m", &[1.0, 0.0], Metric::Cosine)
2189 .unwrap();
2190 txn.commit().unwrap();
2191 }
2192 let mut ro = db.begin(TxnMode::ReadOnly).unwrap();
2193 let err = ro
2194 .search_similar(&[1.0, 0.0, 1.0], Metric::Cosine, 1, None)
2195 .unwrap_err();
2196 assert!(matches!(
2197 err,
2198 Error::Core(alopex_core::Error::DimensionMismatch { .. })
2199 ));
2200 }
2201}