laurus 0.10.0

Unified search library for lexical, vector, and semantic retrieval
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

//! Shared test helpers for integration tests (Issue #889 PR-4).
//!
//! `tests/common/mod.rs` (not `tests/common.rs`) is the standard Cargo idiom
//! for code shared between integration test binaries without being picked
//! up as its own test target.

use std::io::Write;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use laurus::storage::{Storage, StorageInput, StorageOutput};

/// Storage decorator counting every byte written through `create_output`
/// (including rewrites — file sizes alone cannot see them).
///
/// Originated in the #634 HNSW segment-per-commit campaign
/// (`vector_segmented_index_test.rs`); moved here in #889 PR-4 so Flat's
/// (and IVF's, in PR-6) segmented gates can reuse it instead of copying it
/// a third time.
#[derive(Debug)]
pub struct ByteCountingStorage {
    inner: Arc<dyn Storage>,
    pub written: Arc<AtomicU64>,
}

impl ByteCountingStorage {
    pub fn new(inner: Arc<dyn Storage>) -> Self {
        Self {
            inner,
            written: Arc::new(AtomicU64::new(0)),
        }
    }
}

#[derive(Debug)]
struct CountingOutput {
    inner: Box<dyn StorageOutput>,
    written: Arc<AtomicU64>,
}

impl Write for CountingOutput {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let n = self.inner.write(buf)?;
        self.written.fetch_add(n as u64, Ordering::Relaxed);
        Ok(n)
    }
    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

impl std::io::Seek for CountingOutput {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
        self.inner.seek(pos)
    }
}

impl StorageOutput for CountingOutput {
    fn flush_and_sync(&mut self) -> laurus::Result<()> {
        self.inner.flush_and_sync()
    }
    fn position(&self) -> laurus::Result<u64> {
        self.inner.position()
    }
    fn close(&mut self) -> laurus::Result<()> {
        self.inner.close()
    }
}

impl Storage for ByteCountingStorage {
    fn open_input(&self, name: &str) -> laurus::Result<Box<dyn StorageInput>> {
        self.inner.open_input(name)
    }
    fn create_output(&self, name: &str) -> laurus::Result<Box<dyn StorageOutput>> {
        Ok(Box::new(CountingOutput {
            inner: self.inner.create_output(name)?,
            written: self.written.clone(),
        }))
    }
    fn create_output_append(&self, name: &str) -> laurus::Result<Box<dyn StorageOutput>> {
        Ok(Box::new(CountingOutput {
            inner: self.inner.create_output_append(name)?,
            written: self.written.clone(),
        }))
    }
    fn delete_file(&self, name: &str) -> laurus::Result<()> {
        self.inner.delete_file(name)
    }
    fn file_exists(&self, name: &str) -> bool {
        self.inner.file_exists(name)
    }
    fn list_files(&self) -> laurus::Result<Vec<String>> {
        self.inner.list_files()
    }
    fn file_size(&self, name: &str) -> laurus::Result<u64> {
        self.inner.file_size(name)
    }
    fn rename_file(&self, from: &str, to: &str) -> laurus::Result<()> {
        self.inner.rename_file(from, to)
    }
    fn metadata(&self, name: &str) -> laurus::Result<laurus::storage::FileMetadata> {
        self.inner.metadata(name)
    }
    fn create_temp_output(&self, prefix: &str) -> laurus::Result<(String, Box<dyn StorageOutput>)> {
        let (name, output) = self.inner.create_temp_output(prefix)?;
        Ok((
            name,
            Box::new(CountingOutput {
                inner: output,
                written: self.written.clone(),
            }),
        ))
    }
    fn sync(&self) -> laurus::Result<()> {
        self.inner.sync()
    }
    fn close(&mut self) -> laurus::Result<()> {
        Ok(())
    }
}