fcb_core/static_btree/query/
memory.rs

1use chrono::{DateTime, Utc};
2use ordered_float::OrderedFloat;
3use std::collections::HashMap;
4use std::io::{Read, Write};
5
6use crate::static_btree::entry::Entry;
7use crate::static_btree::error::{Error, Result};
8use crate::static_btree::key::{FixedStringKey, Key, KeyType};
9use crate::static_btree::query::types::{Operator, SearchIndex};
10use crate::static_btree::stree::Stree;
11
12use super::types::QueryCondition;
13use super::MultiIndex;
14
15/// In-memory index implementation that wraps the Stree structure
16// NOTE: This can be type alias for Stree later
17#[derive(Debug, Clone)]
18pub struct MemoryIndex<K: Key> {
19    /// The underlying static B-tree
20    stree: Stree<K>,
21}
22
23impl<K: Key> MemoryIndex<K> {
24    /// Create a new memory index from an existing Stree
25    pub fn new(mut data: impl Read, num_items: usize, branching_factor: u16) -> Result<Self> {
26        let stree = Stree::from_buf(&mut data, num_items, branching_factor)?;
27
28        Ok(Self { stree })
29    }
30
31    /// Build a memory index from a collection of entries
32    pub fn build(entries: &[Entry<K>], branching_factor: u16) -> Result<Self> {
33        let stree = Stree::<K>::build(entries, branching_factor)?;
34
35        Ok(Self { stree })
36    }
37
38    pub fn from_buf(mut data: impl Read, num_items: usize, branching_factor: u16) -> Result<Self> {
39        let stree = Stree::from_buf(&mut data, num_items, branching_factor)?;
40
41        Ok(Self { stree })
42    }
43
44    pub fn num_items(&self) -> usize {
45        self.stree.num_leaf_items()
46    }
47
48    pub fn branching_factor(&self) -> u16 {
49        self.stree.branching_factor()
50    }
51
52    pub fn size(&self) -> usize {
53        Stree::<K>::tree_size(self.num_items())
54    }
55
56    pub fn serialize(&self, out: &mut impl Write) -> Result<usize> {
57        self.stree.stream_write(out)
58    }
59
60    pub fn payload_size(&self) -> usize {
61        self.stree.payload_size()
62    }
63
64    /// Find all items in the range, with each bound independently strict
65    /// (exclusive) or inclusive. A `None` bound is the type's min/max sentinel
66    /// and is never strict.
67    ///
68    /// `Gt`/`Lt`/`Ne` lower to this instead of subtracting `find_exact` from
69    /// an inclusive range: the subtraction removes feature offsets, and one
70    /// feature can be indexed under several keys, so it deletes features that
71    /// match through a different key.
72    pub fn find_range_strict(
73        &self,
74        start: Option<K>,
75        start_strict: bool,
76        end: Option<K>,
77        end_strict: bool,
78    ) -> Result<Vec<u64>> {
79        let lower = start.unwrap_or_else(K::min_value);
80        let upper = end.unwrap_or_else(K::max_value);
81        let results = self
82            .stree
83            .find_range_strict(lower, start_strict, upper, end_strict)?;
84        Ok(results.into_iter().map(|item| item.offset as u64).collect())
85    }
86}
87
88impl<K: Key> SearchIndex<K> for MemoryIndex<K> {
89    fn find_exact(&self, key: K) -> Result<Vec<u64>> {
90        let results = self.stree.find_exact(key)?;
91        Ok(results.into_iter().map(|item| item.offset as u64).collect())
92    }
93
94    fn find_range(&self, start: Option<K>, end: Option<K>) -> Result<Vec<u64>> {
95        match (start, end) {
96            (Some(start_key), Some(end_key)) => {
97                let results = self.stree.find_range(start_key, end_key)?;
98                Ok(results.into_iter().map(|item| item.offset as u64).collect())
99            }
100            (Some(start_key), None) => {
101                // Find all items >= start_key
102                let results = self.stree.find_range(start_key, K::max_value())?;
103                Ok(results.into_iter().map(|item| item.offset as u64).collect())
104            }
105            (None, Some(end_key)) => {
106                // Find all items <= end_key
107                let results = self.stree.find_range(K::min_value(), end_key)?;
108                Ok(results.into_iter().map(|item| item.offset as u64).collect())
109            }
110            (None, None) => Err(Error::QueryError(
111                "find_range requires at least one bound".to_string(),
112            )),
113        }
114    }
115}
116
117/// Trait for different index types we might store
118pub trait TypedSearchIndex: Send + Sync {
119    /// Execute the query condition
120    fn execute_query_condition(&self, condition: &QueryCondition) -> Result<Vec<u64>>;
121}
122
123// Macro to implement TypedSearchIndex for each key type following the same pattern
124macro_rules! impl_typed_search_index {
125    ($key_type:ty, $enum_variant:path) => {
126        impl TypedSearchIndex for MemoryIndex<$key_type> {
127            fn execute_query_condition(&self, condition: &QueryCondition) -> Result<Vec<u64>> {
128                // Extract the key value from the enum variant
129                let key = match &condition.key {
130                    $enum_variant(val) => val.clone(),
131                    _ => {
132                        return Err(Error::QueryError(format!(
133                            "key type mismatch: expected {}, got {:?}",
134                            stringify!($key_type),
135                            condition.key
136                        )))
137                    }
138                };
139
140                // Execute query based on operator
141                match condition.operator {
142                    Operator::Eq => self.find_exact(key),
143                    // Two half-open scans rather than a full scan minus the
144                    // equal set: subtraction on feature offsets is wrong when
145                    // one feature carries several values of the attribute.
146                    Operator::Ne => {
147                        let mut results =
148                            self.find_range_strict(None, false, Some(key.clone()), true)?;
149                        let above = self.find_range_strict(Some(key), true, None, false)?;
150                        results.extend(above);
151                        Ok(results)
152                    }
153                    Operator::Gt => self.find_range_strict(Some(key), true, None, false),
154                    Operator::Lt => self.find_range_strict(None, false, Some(key), true),
155                    Operator::Ge => self.find_range_strict(Some(key), false, None, false),
156                    Operator::Le => self.find_range_strict(None, false, Some(key), false),
157                }
158            }
159        }
160    };
161}
162
163// Implement TypedSearchIndex for all supported key types
164impl_typed_search_index!(i32, KeyType::Int32);
165impl_typed_search_index!(i64, KeyType::Int64);
166impl_typed_search_index!(i8, KeyType::Int8);
167impl_typed_search_index!(u8, KeyType::UInt8);
168impl_typed_search_index!(i16, KeyType::Int16);
169impl_typed_search_index!(u16, KeyType::UInt16);
170impl_typed_search_index!(u32, KeyType::UInt32);
171impl_typed_search_index!(u64, KeyType::UInt64);
172impl_typed_search_index!(OrderedFloat<f32>, KeyType::Float32);
173impl_typed_search_index!(OrderedFloat<f64>, KeyType::Float64);
174impl_typed_search_index!(bool, KeyType::Bool);
175impl_typed_search_index!(DateTime<Utc>, KeyType::DateTime);
176impl_typed_search_index!(FixedStringKey<20>, KeyType::StringKey20);
177impl_typed_search_index!(FixedStringKey<50>, KeyType::StringKey50);
178impl_typed_search_index!(FixedStringKey<100>, KeyType::StringKey100);
179
180/// Container for multiple in-memory indices with different key types
181pub struct MemoryMultiIndex {
182    /// Map of field names to typed indices
183    indices: HashMap<String, Box<dyn TypedSearchIndex>>,
184}
185
186impl MemoryMultiIndex {
187    /// Create a new empty multi-index
188    pub fn new() -> Self {
189        Self {
190            indices: HashMap::new(),
191        }
192    }
193
194    /// Generic method to add an index for any supported key type
195    pub fn add_index<K: Key + 'static>(&mut self, field: String, index: MemoryIndex<K>)
196    where
197        MemoryIndex<K>: TypedSearchIndex,
198    {
199        self.indices.insert(field, Box::new(index));
200    }
201
202    pub fn indices(&self) -> &HashMap<String, Box<dyn TypedSearchIndex>> {
203        &self.indices
204    }
205
206    /// Add a string index with key size 20
207    pub fn add_string_index20(&mut self, field: String, index: MemoryIndex<FixedStringKey<20>>) {
208        self.indices.insert(field, Box::new(index));
209    }
210
211    /// Add a string index with key size 50
212    pub fn add_string_index50(&mut self, field: String, index: MemoryIndex<FixedStringKey<50>>) {
213        self.indices.insert(field, Box::new(index));
214    }
215
216    /// Add a string index with key size 100
217    pub fn add_string_index100(&mut self, field: String, index: MemoryIndex<FixedStringKey<100>>) {
218        self.indices.insert(field, Box::new(index));
219    }
220
221    /// Add an i32 index
222    pub fn add_i32_index(&mut self, field: String, index: MemoryIndex<i32>) {
223        self.indices.insert(field, Box::new(index));
224    }
225
226    /// Add an i64 index
227    pub fn add_i64_index(&mut self, field: String, index: MemoryIndex<i64>) {
228        self.indices.insert(field, Box::new(index));
229    }
230
231    /// Add a u32 index
232    pub fn add_u32_index(&mut self, field: String, index: MemoryIndex<u32>) {
233        self.indices.insert(field, Box::new(index));
234    }
235
236    /// Add a u64 index
237    pub fn add_u64_index(&mut self, field: String, index: MemoryIndex<u64>) {
238        self.indices.insert(field, Box::new(index));
239    }
240
241    /// Add a float32 index
242    pub fn add_f32_index(&mut self, field: String, index: MemoryIndex<OrderedFloat<f32>>) {
243        self.indices.insert(field, Box::new(index));
244    }
245
246    /// Add a float64 index
247    pub fn add_f64_index(&mut self, field: String, index: MemoryIndex<OrderedFloat<f64>>) {
248        self.indices.insert(field, Box::new(index));
249    }
250
251    /// Add a i8 index
252    pub fn add_i8_index(&mut self, field: String, index: MemoryIndex<i8>) {
253        self.indices.insert(field, Box::new(index));
254    }
255
256    /// Add a u8 index
257    pub fn add_u8_index(&mut self, field: String, index: MemoryIndex<u8>) {
258        self.indices.insert(field, Box::new(index));
259    }
260
261    /// Add a i16 index
262    pub fn add_i16_index(&mut self, field: String, index: MemoryIndex<i16>) {
263        self.indices.insert(field, Box::new(index));
264    }
265
266    /// Add a u16 index
267    pub fn add_u16_index(&mut self, field: String, index: MemoryIndex<u16>) {
268        self.indices.insert(field, Box::new(index));
269    }
270
271    /// Add a boolean index
272    pub fn add_bool_index(&mut self, field: String, index: MemoryIndex<bool>) {
273        self.indices.insert(field, Box::new(index));
274    }
275
276    /// Add a datetime index
277    pub fn add_datetime_index(&mut self, field: String, index: MemoryIndex<DateTime<Utc>>) {
278        self.indices.insert(field, Box::new(index));
279    }
280}
281
282impl MultiIndex for MemoryMultiIndex {
283    /// Execute a heterogeneous query with different key types
284    fn query(&self, conditions: &[QueryCondition]) -> Result<Vec<u64>> {
285        if conditions.is_empty() {
286            return Err(Error::QueryError("query cannot be empty".to_string()));
287        }
288
289        // Process the first condition to initialize the result set
290        let first_condition = &conditions[0];
291        let index = self.indices.get(&first_condition.field).ok_or_else(|| {
292            Error::QueryError(format!(
293                "no index found for field '{}'",
294                first_condition.field
295            ))
296        })?;
297        let mut result_set = index.execute_query_condition(first_condition)?;
298        if result_set.is_empty() {
299            return Ok(vec![]);
300        }
301
302        // Process remaining conditions with set intersection
303        for condition in &conditions[1..] {
304            let index = self.indices.get(&condition.field).ok_or_else(|| {
305                Error::QueryError(format!("no index found for field '{}'", condition.field))
306            })?;
307            let condition_results = index.execute_query_condition(condition)?;
308
309            // Perform intersection (AND logic)
310            result_set.retain(|offset| condition_results.contains(offset));
311
312            // If result set is empty, we can short-circuit. For now it's AND logic, so if any condition is empty, the result set is empty
313            if result_set.is_empty() {
314                return Ok(vec![]);
315            }
316        }
317
318        Ok(result_set)
319    }
320}
321
322impl Default for MemoryMultiIndex {
323    fn default() -> Self {
324        Self::new()
325    }
326}
fcb_core/static_btree/query/memory.rs

fcb_core/static_btree/query/
memory.rs
