fcb_core/static_btree/query/
stream.rs

1use std::collections::HashMap;
2use std::fmt::Debug;
3use std::io::{Read, Seek, SeekFrom};
4use std::marker::PhantomData;
5use std::ops::Range;
6
7use chrono::{DateTime, Utc};
8use ordered_float::OrderedFloat;
9
10use crate::static_btree::error::{Error, Result};
11use crate::static_btree::key::{FixedStringKey, Key, KeyType};
12use crate::static_btree::query::types::{Operator, QueryCondition};
13use crate::static_btree::stree::Stree;
14
15/// Stream-based index for file access
16#[derive(Debug, Clone)]
17pub struct StreamIndex<K: Key> {
18    /// Number of items in the index
19    num_items: usize,
20    /// Branching factor of the tree
21    branching_factor: u16,
22    /// Offset of the index in the file
23    index_offset: u64,
24    /// Size of the index
25    length: u64,
26    /// Phantom marker for the key type
27    _marker: PhantomData<K>,
28}
29
30impl<K: Key> StreamIndex<K> {
31    /// Create a new stream index with metadata
32    pub fn new(num_items: usize, branching_factor: u16, index_offset: u64, length: u64) -> Self {
33        Self {
34            num_items,
35            branching_factor,
36            index_offset,
37            length,
38            _marker: PhantomData,
39        }
40    }
41
42    /// Get the number of items in the index
43    pub fn num_items(&self) -> usize {
44        self.num_items
45    }
46
47    /// Get the branching factor of the tree
48    pub fn branching_factor(&self) -> u16 {
49        self.branching_factor
50    }
51
52    /// Get the index offset
53    pub fn index_offset(&self) -> u64 {
54        self.index_offset
55    }
56
57    /// Get the length of the index
58    pub fn length(&self) -> u64 {
59        self.length
60    }
61
62    /// Find exact matches using a reader
63    pub fn find_exact_with_reader<R: Read + Seek + ?Sized>(
64        &self,
65        reader: &mut R,
66        key: K,
67    ) -> Result<Vec<u64>> {
68        let results = Stree::stream_find_exact(reader, self.num_items, self.branching_factor, key)?;
69
70        Ok(results.into_iter().map(|item| item.offset as u64).collect())
71    }
72
73    /// Find all items in the range using a reader, with each bound
74    /// independently strict (exclusive) or inclusive. A `None` bound is the
75    /// type's min/max sentinel and is never strict.
76    ///
77    /// `Gt`/`Lt`/`Ne` lower to this instead of subtracting `find_exact` from
78    /// an inclusive range: the subtraction removes feature offsets, and one
79    /// feature can be indexed under several keys, so it deletes features that
80    /// match through a different key.
81    pub fn find_range_strict_with_reader<R: Read + Seek + ?Sized>(
82        &self,
83        reader: &mut R,
84        start: Option<K>,
85        start_strict: bool,
86        end: Option<K>,
87        end_strict: bool,
88    ) -> Result<Vec<u64>> {
89        let start_position = reader.stream_position()?;
90        let lower = start.unwrap_or_else(K::min_value);
91        let upper = end.unwrap_or_else(K::max_value);
92        let results = Stree::stream_find_range_strict(
93            reader,
94            self.num_items,
95            self.branching_factor,
96            lower,
97            start_strict,
98            upper,
99            end_strict,
100        );
101        reader.seek(SeekFrom::Start(start_position))?;
102        Ok(results?
103            .into_iter()
104            .map(|item| item.offset as u64)
105            .collect())
106    }
107
108    /// Find range matches using a reader
109    pub fn find_range_with_reader<R: Read + Seek + ?Sized>(
110        &self,
111        reader: &mut R,
112        start: Option<K>,
113        end: Option<K>,
114    ) -> Result<Vec<u64>> {
115        // print current cursor position
116        let start_position = reader.stream_position()?;
117        let results = match (start, end) {
118            (Some(start_key), Some(end_key)) => {
119                let results = Stree::stream_find_range(
120                    reader,
121                    self.num_items,
122                    self.branching_factor,
123                    start_key,
124                    end_key,
125                )?;
126                Ok(results.into_iter().map(|item| item.offset as u64).collect())
127            }
128            (Some(start_key), None) => {
129                // Find all items >= start_key
130                let results = Stree::stream_find_range(
131                    reader,
132                    self.num_items,
133                    self.branching_factor,
134                    start_key,
135                    K::max_value(),
136                )?;
137                Ok(results.into_iter().map(|item| item.offset as u64).collect())
138            }
139            (None, Some(end_key)) => {
140                // Find all items <= end_key
141                let results = Stree::stream_find_range(
142                    reader,
143                    self.num_items,
144                    self.branching_factor,
145                    K::min_value(),
146                    end_key,
147                )?;
148                Ok(results.into_iter().map(|item| item.offset as u64).collect())
149            }
150            (None, None) => Err(Error::QueryError(
151                "find_range requires at least one bound".to_string(),
152            )),
153        };
154
155        reader.seek(SeekFrom::Start(start_position))?;
156        results
157    }
158}
159
160/// Trait alias for objects that implement Read and Seek, to allow trait objects
161pub trait ReadSeek: Read + Seek {}
162impl<T: Read + Seek> ReadSeek for T {}
163
164/// Trait for typed stream search index with heterogeneous key types
165pub trait TypedStreamSearchIndex: Send + Sync {
166    /// Execute the query condition using the provided reader
167    fn execute_query_condition(
168        &self,
169        reader: &mut dyn ReadSeek,
170        condition: &QueryCondition,
171    ) -> Result<Vec<u64>>;
172}
173
174// Macro to implement TypedStreamSearchIndex for each supported key type
175macro_rules! impl_typed_stream_search_index {
176    ($key_type:ty, $enum_variant:path) => {
177        impl TypedStreamSearchIndex for StreamIndex<$key_type> {
178            fn execute_query_condition(
179                &self,
180                reader: &mut dyn ReadSeek,
181                condition: &QueryCondition,
182            ) -> Result<Vec<u64>> {
183                let start_position = reader.stream_position()?;
184                // Extract the key value from the enum variant
185                let key = match &condition.key {
186                    $enum_variant(val) => val.clone(),
187                    _ => {
188                        return Err(Error::QueryError(format!(
189                            "key type mismatch: expected {}, got {:?}",
190                            stringify!($key_type),
191                            condition.key
192                        )))
193                    }
194                };
195                // Execute query based on operator
196                let items = match condition.operator {
197                    Operator::Eq => self.find_exact_with_reader(reader, key)?,
198                    // Two half-open scans rather than a full scan minus the
199                    // equal set: subtraction on feature offsets is wrong when
200                    // one feature carries several values of the attribute.
201                    Operator::Ne => {
202                        let mut results = self.find_range_strict_with_reader(
203                            reader,
204                            None,
205                            false,
206                            Some(key.clone()),
207                            true,
208                        )?;
209                        let above = self.find_range_strict_with_reader(
210                            reader,
211                            Some(key),
212                            true,
213                            None,
214                            false,
215                        )?;
216                        results.extend(above);
217                        results
218                    }
219                    Operator::Gt => {
220                        self.find_range_strict_with_reader(reader, Some(key), true, None, false)?
221                    }
222                    Operator::Lt => {
223                        self.find_range_strict_with_reader(reader, None, false, Some(key), true)?
224                    }
225                    Operator::Ge => {
226                        self.find_range_strict_with_reader(reader, Some(key), false, None, false)?
227                    }
228                    Operator::Le => {
229                        self.find_range_strict_with_reader(reader, None, false, Some(key), false)?
230                    }
231                };
232                reader.seek(SeekFrom::Start(start_position))?;
233                Ok(items)
234            }
235        }
236    };
237}
238
239// Implement TypedStreamSearchIndex for all supported key types
240impl_typed_stream_search_index!(i8, KeyType::Int8);
241impl_typed_stream_search_index!(u8, KeyType::UInt8);
242impl_typed_stream_search_index!(i16, KeyType::Int16);
243impl_typed_stream_search_index!(u16, KeyType::UInt16);
244impl_typed_stream_search_index!(i32, KeyType::Int32);
245impl_typed_stream_search_index!(i64, KeyType::Int64);
246impl_typed_stream_search_index!(u32, KeyType::UInt32);
247impl_typed_stream_search_index!(u64, KeyType::UInt64);
248impl_typed_stream_search_index!(OrderedFloat<f32>, KeyType::Float32);
249impl_typed_stream_search_index!(OrderedFloat<f64>, KeyType::Float64);
250impl_typed_stream_search_index!(bool, KeyType::Bool);
251impl_typed_stream_search_index!(DateTime<Utc>, KeyType::DateTime);
252impl_typed_stream_search_index!(FixedStringKey<20>, KeyType::StringKey20);
253impl_typed_stream_search_index!(FixedStringKey<50>, KeyType::StringKey50);
254impl_typed_stream_search_index!(FixedStringKey<100>, KeyType::StringKey100);
255
256/// Container for multiple stream indices with different key types
257pub struct StreamMultiIndex {
258    indices: HashMap<String, Box<dyn TypedStreamSearchIndex>>,
259    index_offsets: HashMap<String, Range<usize>>,
260}
261
262impl StreamMultiIndex {
263    /// Create a new empty multi-index
264    pub fn new() -> Self {
265        Self {
266            indices: HashMap::new(),
267            index_offsets: HashMap::new(),
268        }
269    }
270
271    /// Generic method to add an index for any supported key type
272    pub fn add_index<K: Key + 'static>(&mut self, field: String, index: StreamIndex<K>)
273    where
274        StreamIndex<K>: TypedStreamSearchIndex,
275    {
276        self.indices.insert(field, Box::new(index));
277    }
278
279    fn add_index_offset(&mut self, field: String, length: u64) {
280        //length of the index about to be added
281        // get the last index offset
282        let largest_offset = self
283            .index_offsets
284            .values()
285            .map(|v| v.end)
286            .max()
287            .unwrap_or(0);
288        self.index_offsets
289            .insert(field, largest_offset..largest_offset + length as usize);
290    }
291
292    /// Add a string index with key size 20
293    pub fn add_string_index20(
294        &mut self,
295        field: String,
296        index: StreamIndex<FixedStringKey<20>>,
297        length: u64,
298    ) {
299        self.indices.insert(field.clone(), Box::new(index));
300        self.add_index_offset(field, length);
301    }
302
303    /// Add a string index with key size 50
304    pub fn add_string_index50(
305        &mut self,
306        field: String,
307        index: StreamIndex<FixedStringKey<50>>,
308        length: u64,
309    ) {
310        self.indices.insert(field.clone(), Box::new(index));
311        self.add_index_offset(field, length);
312    }
313
314    /// Add a string index with key size 100
315    pub fn add_string_index100(
316        &mut self,
317        field: String,
318        index: StreamIndex<FixedStringKey<100>>,
319        length: u64,
320    ) {
321        self.indices.insert(field.clone(), Box::new(index));
322        self.add_index_offset(field, length);
323    }
324
325    /// Add an i8 index
326    pub fn add_i8_index(&mut self, field: String, index: StreamIndex<i8>, length: u64) {
327        self.indices.insert(field.clone(), Box::new(index));
328        self.add_index_offset(field, length);
329    }
330
331    /// Add a u8 index
332    pub fn add_u8_index(&mut self, field: String, index: StreamIndex<u8>, length: u64) {
333        self.indices.insert(field.clone(), Box::new(index));
334        self.add_index_offset(field, length);
335    }
336
337    /// Add an i16 index
338    pub fn add_i16_index(&mut self, field: String, index: StreamIndex<i16>, length: u64) {
339        self.indices.insert(field.clone(), Box::new(index));
340        self.add_index_offset(field, length);
341    }
342
343    /// Add a u16 index
344    pub fn add_u16_index(&mut self, field: String, index: StreamIndex<u16>, length: u64) {
345        self.indices.insert(field.clone(), Box::new(index));
346        self.add_index_offset(field, length);
347    }
348
349    /// Add an i32 index
350    pub fn add_i32_index(&mut self, field: String, index: StreamIndex<i32>, length: u64) {
351        self.indices.insert(field.clone(), Box::new(index));
352        self.add_index_offset(field, length);
353    }
354
355    /// Add an i64 index
356    pub fn add_i64_index(&mut self, field: String, index: StreamIndex<i64>, length: u64) {
357        self.indices.insert(field.clone(), Box::new(index));
358        self.add_index_offset(field, length);
359    }
360
361    /// Add a u32 index
362    pub fn add_u32_index(&mut self, field: String, index: StreamIndex<u32>, length: u64) {
363        self.indices.insert(field.clone(), Box::new(index));
364        self.add_index_offset(field, length);
365    }
366
367    /// Add a u64 index
368    pub fn add_u64_index(&mut self, field: String, index: StreamIndex<u64>, length: u64) {
369        self.indices.insert(field.clone(), Box::new(index));
370        self.add_index_offset(field, length);
371    }
372
373    /// Add a float32 index
374    pub fn add_f32_index(
375        &mut self,
376        field: String,
377        index: StreamIndex<OrderedFloat<f32>>,
378        length: u64,
379    ) {
380        self.indices.insert(field.clone(), Box::new(index));
381        self.add_index_offset(field, length);
382    }
383
384    /// Add a float64 index
385    pub fn add_f64_index(
386        &mut self,
387        field: String,
388        index: StreamIndex<OrderedFloat<f64>>,
389        length: u64,
390    ) {
391        self.indices.insert(field.clone(), Box::new(index));
392        self.add_index_offset(field, length);
393    }
394
395    /// Add a boolean index
396    pub fn add_bool_index(&mut self, field: String, index: StreamIndex<bool>, length: u64) {
397        self.indices.insert(field.clone(), Box::new(index));
398        self.add_index_offset(field, length);
399    }
400
401    /// Add a datetime index
402    pub fn add_datetime_index(
403        &mut self,
404        field: String,
405        index: StreamIndex<DateTime<Utc>>,
406        length: u64,
407    ) {
408        self.indices.insert(field.clone(), Box::new(index));
409        self.add_index_offset(field, length);
410    }
411
412    /// Execute a heterogeneous query with different key types using a reader
413    pub fn query(
414        &self,
415        reader: &mut dyn ReadSeek,
416        conditions: &[QueryCondition],
417    ) -> Result<Vec<u64>> {
418        if conditions.is_empty() {
419            return Err(Error::QueryError("query cannot be empty".to_string()));
420        }
421        let first = &conditions[0];
422        let indexer = self.indices.get(&first.field).ok_or_else(|| {
423            Error::QueryError(format!("no index found for field '{}'", first.field))
424        })?;
425        let index_range = self.index_offsets.get(&first.field).ok_or_else(|| {
426            Error::QueryError(format!("no index range found for field '{}'", first.field))
427        })?;
428
429        // currently reader is continuous buffer of multiple indices. We need to create different readers for each index. `index_offsets` field of the struct accomodates Range of each indices. e.g. if index_offsets is [(field1, 0..100), (field2, 100..200)], it means that field1 is at offset 0-99 and field2 is at offset 100-199 in the reader. Since `execute_query_condition` is called with a reader, we need to create a new reader for each index.
430
431        let start_position = reader.stream_position()?;
432        // set cursor to the start of the index
433        reader.seek(SeekFrom::Start(start_position + index_range.start as u64))?;
434
435        let mut result_set = indexer.execute_query_condition(reader, first)?;
436        // Restore the cursor BEFORE the early return: leaving it parked on
437        // this index's start would make the next query resolve every index
438        // offset relative to the wrong base.
439        reader.seek(SeekFrom::Start(start_position))?;
440        if result_set.is_empty() {
441            return Ok(vec![]);
442        }
443
444        for cond in &conditions[1..] {
445            let start_position = reader.stream_position()?;
446            let indexer = self.indices.get(&cond.field).ok_or_else(|| {
447                Error::QueryError(format!("no index found for field '{}'", cond.field))
448            })?;
449            let index_range = self.index_offsets.get(&cond.field).ok_or_else(|| {
450                Error::QueryError(format!("no index range found for field '{}'", cond.field))
451            })?;
452            let index_start = start_position + index_range.start as u64;
453            // set cursor to the start of the index
454            reader.seek(SeekFrom::Start(index_start))?;
455            println!("index_start: {index_start}");
456            println!("start_position: {start_position}");
457            println!("query condition: {cond:?}");
458            let condition_results = indexer.execute_query_condition(reader, cond)?;
459            // set cursor to the start of the index, before any early return
460            reader.seek(SeekFrom::Start(start_position))?;
461            result_set.retain(|offset| condition_results.contains(offset));
462            if result_set.is_empty() {
463                return Ok(vec![]); // no results found for this condition, return early so we don't waste time intersecting empty sets
464            }
465        }
466        // set cursor to the start of the index
467        reader.seek(SeekFrom::Start(start_position))?;
468        Ok(result_set)
469    }
470}
471
472impl Default for StreamMultiIndex {
473    fn default() -> Self {
474        Self::new()
475    }
476}
fcb_core/static_btree/query/stream.rs

fcb_core/static_btree/query/
stream.rs
