fcb_core/reader/
attr_query.rs

1use crate::static_btree::{
2    FixedStringKey, Float, KeyType, MemoryIndex, MemoryMultiIndex, MultiIndex, Operator, Query,
3    QueryCondition, StreamIndex, StreamMultiIndex,
4};
5use std::collections::HashMap;
6use std::io::{self, Cursor, Read, Seek, SeekFrom};
7use std::ops::Range;
8
9use crate::error::{Error, Result};
10
11use chrono::{DateTime, Utc};
12
13use crate::fb::Column;
14use crate::fb::ColumnType;
15use crate::{AttributeIndex, FeatureOffset};
16
17use super::{
18    reader_trait::{NotSeekable, Seekable},
19    FcbReader, FeatureIter,
20};
21
22pub type AttrQuery = Vec<(String, Operator, KeyType)>;
23
24pub fn add_indices_to_multi_memory_index<R: Read>(
25    mut data: R,
26    multi_index: &mut MemoryMultiIndex,
27    columns: &[Column],
28    query: &AttrQuery,
29    attr_info: &AttributeIndex,
30) -> Result<()> {
31    let length = attr_info.length();
32    let mut buf = vec![0; length as usize];
33    data.read_exact(&mut buf)?;
34    let mut buf = Cursor::new(buf);
35    if let Some(col) = columns.iter().find(|col| col.index() == attr_info.index()) {
36        if query.iter().any(|(name, _, _)| col.name() == name) {
37            match col.type_() {
38                ColumnType::Int => {
39                    let index = MemoryIndex::<i32>::from_buf(
40                        &mut buf,
41                        attr_info.num_unique_items() as usize,
42                        attr_info.branching_factor(),
43                    )?;
44                    multi_index.add_i32_index(col.name().to_string(), index);
45                }
46                ColumnType::Float => {
47                    let index = MemoryIndex::<Float<f32>>::from_buf(
48                        &mut buf,
49                        attr_info.num_unique_items() as usize,
50                        attr_info.branching_factor(),
51                    )?;
52                    multi_index.add_f32_index(col.name().to_string(), index);
53                }
54                ColumnType::Double => {
55                    let index = MemoryIndex::<Float<f64>>::from_buf(
56                        &mut buf,
57                        attr_info.num_unique_items() as usize,
58                        attr_info.branching_factor(),
59                    )?;
60                    multi_index.add_f64_index(col.name().to_string(), index);
61                }
62                ColumnType::String => {
63                    let index = MemoryIndex::<FixedStringKey<50>>::from_buf(
64                        &mut buf,
65                        attr_info.num_unique_items() as usize,
66                        attr_info.branching_factor(),
67                    )?;
68                    multi_index.add_string_index50(col.name().to_string(), index);
69                }
70                ColumnType::Bool => {
71                    let index = MemoryIndex::<bool>::from_buf(
72                        &mut buf,
73                        attr_info.num_unique_items() as usize,
74                        attr_info.branching_factor(),
75                    )?;
76                    multi_index.add_bool_index(col.name().to_string(), index);
77                }
78                ColumnType::DateTime => {
79                    let index = MemoryIndex::<DateTime<Utc>>::from_buf(
80                        &mut buf,
81                        attr_info.num_unique_items() as usize,
82                        attr_info.branching_factor(),
83                    )?;
84                    multi_index.add_datetime_index(col.name().to_string(), index);
85                }
86                ColumnType::Short => {
87                    let index = MemoryIndex::<i16>::from_buf(
88                        &mut buf,
89                        attr_info.num_unique_items() as usize,
90                        attr_info.branching_factor(),
91                    )?;
92                    multi_index.add_i16_index(col.name().to_string(), index);
93                }
94                ColumnType::UShort => {
95                    let index = MemoryIndex::<u16>::from_buf(
96                        &mut buf,
97                        attr_info.num_unique_items() as usize,
98                        attr_info.branching_factor(),
99                    )?;
100                    multi_index.add_u16_index(col.name().to_string(), index);
101                }
102                ColumnType::UInt => {
103                    let index = MemoryIndex::<u32>::from_buf(
104                        &mut buf,
105                        attr_info.num_unique_items() as usize,
106                        attr_info.branching_factor(),
107                    )?;
108                    multi_index.add_u32_index(col.name().to_string(), index);
109                }
110                ColumnType::Long => {
111                    let index = MemoryIndex::<i64>::from_buf(
112                        &mut buf,
113                        attr_info.num_unique_items() as usize,
114                        attr_info.branching_factor(),
115                    )?;
116                    multi_index.add_i64_index(col.name().to_string(), index);
117                }
118                ColumnType::ULong => {
119                    let index = MemoryIndex::<u64>::from_buf(
120                        &mut buf,
121                        attr_info.num_unique_items() as usize,
122                        attr_info.branching_factor(),
123                    )?;
124                    multi_index.add_u64_index(col.name().to_string(), index);
125                }
126                // Byte is stored as u8 by the writer (writer/attribute.rs)
127                // and indexed as MemoryIndex<u8> (writer/attr_index.rs), so it
128                // must be read back as u8. Decoding it as i8 turned every
129                // stored value above 127 into a negative number that was never
130                // written.
131                ColumnType::Byte => {
132                    let index = MemoryIndex::<u8>::from_buf(
133                        &mut buf,
134                        attr_info.num_unique_items() as usize,
135                        attr_info.branching_factor(),
136                    )?;
137                    multi_index.add_u8_index(col.name().to_string(), index);
138                }
139                ColumnType::UByte => {
140                    let index = MemoryIndex::<u8>::from_buf(
141                        &mut buf,
142                        attr_info.num_unique_items() as usize,
143                        attr_info.branching_factor(),
144                    )?;
145                    multi_index.add_u8_index(col.name().to_string(), index);
146                }
147                _ => return Err(Error::UnsupportedColumnType(col.name().to_string())),
148            }
149        } else {
150            println!("  - Skipping index for field: {}", col.name());
151        }
152    }
153    Ok(())
154}
155
156pub fn add_indices_to_multi_stream_index<R: Read + Seek>(
157    multi_index: &mut StreamMultiIndex,
158    columns: &[Column],
159    attr_info: &AttributeIndex,
160    index_begin: usize,
161) -> Result<()> {
162    if let Some(col) = columns.iter().find(|col| col.index() == attr_info.index()) {
163        // TODO: now it assuming to add all indices to the multi_index. However, we should only add the indices that are used in the query. To do that, we need to change the implementation of StreamMultiIndex. Current StreamMultiIndex's `add_index` method assumes that all indices are added to the multi_index. We'll change it to take Range<usize> as an argument.
164        let index_begin = index_begin as u64;
165        match col.type_() {
166            ColumnType::Int => {
167                let index = StreamIndex::<i32>::new(
168                    attr_info.num_unique_items() as usize,
169                    attr_info.branching_factor(),
170                    index_begin,
171                    attr_info.length() as u64,
172                );
173                multi_index.add_i32_index(col.name().to_string(), index, attr_info.length() as u64);
174            }
175            ColumnType::Float => {
176                let index = StreamIndex::<Float<f32>>::new(
177                    attr_info.num_unique_items() as usize,
178                    attr_info.branching_factor(),
179                    index_begin,
180                    attr_info.length() as u64,
181                );
182                multi_index.add_f32_index(col.name().to_string(), index, attr_info.length() as u64);
183            }
184            ColumnType::Double => {
185                let index = StreamIndex::<Float<f64>>::new(
186                    attr_info.num_unique_items() as usize,
187                    attr_info.branching_factor(),
188                    index_begin,
189                    attr_info.length() as u64,
190                );
191                multi_index.add_f64_index(col.name().to_string(), index, attr_info.length() as u64);
192            }
193            ColumnType::String => {
194                let index = StreamIndex::<FixedStringKey<50>>::new(
195                    attr_info.num_unique_items() as usize,
196                    attr_info.branching_factor(),
197                    index_begin,
198                    attr_info.length() as u64,
199                );
200                multi_index.add_string_index50(
201                    col.name().to_string(),
202                    index,
203                    attr_info.length() as u64,
204                );
205            }
206            ColumnType::Bool => {
207                let index = StreamIndex::<bool>::new(
208                    attr_info.num_unique_items() as usize,
209                    attr_info.branching_factor(),
210                    index_begin,
211                    attr_info.length() as u64,
212                );
213                multi_index.add_bool_index(
214                    col.name().to_string(),
215                    index,
216                    attr_info.length() as u64,
217                );
218            }
219            ColumnType::DateTime => {
220                let index = StreamIndex::<DateTime<Utc>>::new(
221                    attr_info.num_unique_items() as usize,
222                    attr_info.branching_factor(),
223                    index_begin,
224                    attr_info.length() as u64,
225                );
226                multi_index.add_datetime_index(
227                    col.name().to_string(),
228                    index,
229                    attr_info.length() as u64,
230                );
231            }
232            ColumnType::Short => {
233                let index = StreamIndex::<i16>::new(
234                    attr_info.num_unique_items() as usize,
235                    attr_info.branching_factor(),
236                    index_begin,
237                    attr_info.length() as u64,
238                );
239                multi_index.add_i16_index(col.name().to_string(), index, attr_info.length() as u64);
240            }
241            ColumnType::UShort => {
242                let index = StreamIndex::<u16>::new(
243                    attr_info.num_unique_items() as usize,
244                    attr_info.branching_factor(),
245                    index_begin,
246                    attr_info.length() as u64,
247                );
248                multi_index.add_u16_index(col.name().to_string(), index, attr_info.length() as u64);
249            }
250            ColumnType::UInt => {
251                let index = StreamIndex::<u32>::new(
252                    attr_info.num_unique_items() as usize,
253                    attr_info.branching_factor(),
254                    index_begin,
255                    attr_info.length() as u64,
256                );
257                multi_index.add_u32_index(col.name().to_string(), index, attr_info.length() as u64);
258            }
259            ColumnType::Long => {
260                let index = StreamIndex::<i64>::new(
261                    attr_info.num_unique_items() as usize,
262                    attr_info.branching_factor(),
263                    index_begin,
264                    attr_info.length() as u64,
265                );
266                multi_index.add_i64_index(col.name().to_string(), index, attr_info.length() as u64);
267            }
268            ColumnType::ULong => {
269                let index = StreamIndex::<u64>::new(
270                    attr_info.num_unique_items() as usize,
271                    attr_info.branching_factor(),
272                    index_begin,
273                    attr_info.length() as u64,
274                );
275                multi_index.add_u64_index(col.name().to_string(), index, attr_info.length() as u64);
276            }
277            ColumnType::Byte => {
278                // See the Byte note above: the writer stores u8.
279                let index = StreamIndex::<u8>::new(
280                    attr_info.num_unique_items() as usize,
281                    attr_info.branching_factor(),
282                    index_begin,
283                    attr_info.length() as u64,
284                );
285                multi_index.add_u8_index(col.name().to_string(), index, attr_info.length() as u64);
286            }
287            ColumnType::UByte => {
288                let index = StreamIndex::<u8>::new(
289                    attr_info.num_unique_items() as usize,
290                    attr_info.branching_factor(),
291                    index_begin,
292                    attr_info.length() as u64,
293                );
294                multi_index.add_u8_index(col.name().to_string(), index, attr_info.length() as u64);
295            }
296            _ => return Err(Error::UnsupportedColumnType(col.name().to_string())),
297        }
298        // }
299        // else {
300        //     println!("  - Skipping index for field: {}", col.name());
301        // }
302    }
303    Ok(())
304}
305
306pub fn build_query(query: &AttrQuery) -> Query {
307    let conditions = query
308        .iter()
309        .map(|(field, operator, key)| {
310            let owned_key = key.clone();
311            QueryCondition {
312                field: field.clone(),
313                operator: *operator,
314                key: owned_key,
315            }
316        })
317        .collect();
318    Query { conditions }
319}
320
321impl<R: Read + Seek> FcbReader<R> {
322    pub fn select_attr_query(mut self, query: AttrQuery) -> Result<FeatureIter<R, Seekable>> {
323        // query: vec<(field_name, operator, value)>
324        let header = self.buffer.header();
325        let attr_index_entries = header
326            .attribute_index()
327            .ok_or(Error::AttributeIndexNotFound)?;
328        if attr_index_entries.is_empty() {
329            return Err(Error::AttributeIndexNotFound);
330        }
331
332        let mut attr_index_entries: Vec<&AttributeIndex> = attr_index_entries.iter().collect();
333        attr_index_entries.sort_by_key(|attr| attr.index());
334
335        let columns = header
336            .columns()
337            .ok_or(Error::NoColumnsInHeader)?
338            .iter()
339            .collect::<Vec<_>>();
340
341        // Range of attribute indices to be processed. HashMap<field_name, Range<usize>>
342        let mut attr_index_range = HashMap::<String, Range<usize>>::new();
343        let mut current_index = 0;
344        for attr_info in attr_index_entries.iter() {
345            let column = columns
346                .iter()
347                .find(|c| c.index() == attr_info.index())
348                .ok_or(Error::AttributeIndexNotFound)?;
349            let field_name = column.name().to_string();
350            let index_begin = current_index;
351            let index_end = index_begin + attr_info.length() as usize;
352            attr_index_range.insert(
353                field_name,
354                Range {
355                    start: index_begin,
356                    end: index_end,
357                },
358            );
359            current_index = index_end;
360        }
361
362        // Get the current position (should be at the start of the file)
363        // let start_pos = self.reader.stream_position()?;
364
365        // Skip the rtree index bytes; we know the correct offset for that
366        let rtree_offset = self.rtree_index_size();
367        self.reader.seek(SeekFrom::Current(rtree_offset as i64))?;
368
369        // Now we should be at the start of the attribute indices
370        let attr_index_start_pos = self.reader.stream_position()?;
371
372        // Reset reader position to the start of attribute indices
373        self.reader.seek(SeekFrom::Start(attr_index_start_pos))?;
374
375        // Create a query from the AttrQuery
376        let query_obj = build_query(&query);
377
378        let mut multi_index = StreamMultiIndex::new();
379        // iterate over the columens which are used in the query and is in columns and in attr_index_entries
380        for attr_info in attr_index_entries.iter() {
381            let column_idx = attr_info.index();
382            let column = columns
383                .iter()
384                .find(|c| c.index() == column_idx)
385                .ok_or(Error::AttributeIndexNotFound)?;
386            // if query
387            //     .iter()
388            //     .any(|(name, _, _)| name.as_str() == column.name())
389
390            let index_range = attr_index_range
391                .get(column.name())
392                .ok_or(Error::AttributeIndexNotFound)?;
393            add_indices_to_multi_stream_index::<R>(
394                &mut multi_index,
395                &columns,
396                attr_info,
397                index_range.start,
398            )?;
399        }
400
401        let result = match multi_index.query(&mut self.reader, &query_obj.conditions) {
402            Ok(res) => res,
403            Err(e) => {
404                return Err(Error::QueryExecutionError(format!(
405                    "Failed to execute streaming query: {e}"
406                )));
407            }
408        };
409
410        // Sort the results
411        let mut result_vec: Vec<u64> = result.into_iter().collect();
412        result_vec.sort();
413
414        let header_size = self.buffer.header_buf.len();
415        let feature_offset = FeatureOffset {
416            magic_bytes: 8,
417            header: header_size as u64,
418            rtree_index: self.rtree_index_size(),
419            attributes: self.attr_index_size(),
420        };
421
422        let total_feat_count = result_vec.len() as u64;
423
424        let attr_index_size = self.attr_index_size();
425        self.reader
426            .seek(SeekFrom::Start(attr_index_start_pos + attr_index_size))?;
427
428        Ok(FeatureIter::<R, Seekable>::new(
429            self.reader,
430            self.verify,
431            self.buffer,
432            None,
433            Some(result_vec),
434            feature_offset,
435            total_feat_count,
436        ))
437    }
438}
439
440impl<R: Read> FcbReader<R> {
441    pub fn select_attr_query_seq(
442        mut self,
443        query: AttrQuery,
444    ) -> Result<FeatureIter<R, NotSeekable>> {
445        // query: vec<(field_name, operator, value)>
446        let header = self.buffer.header();
447        let attr_index_entries = header
448            .attribute_index()
449            .ok_or(Error::AttributeIndexNotFound)?;
450        let columns: Vec<Column> = header
451            .columns()
452            .ok_or(Error::NoColumnsInHeader)?
453            .iter()
454            .collect();
455
456        // Instead of seeking, read and discard the rtree index bytes; we know the correct offset for that.
457        let rtree_offset = self.rtree_index_size();
458        io::copy(&mut (&mut self.reader).take(rtree_offset), &mut io::sink())?;
459
460        // Since we can't use StreamableMultiIndex with a non-seekable reader,
461        // we'll still use MultiIndex but optimize the process to minimize memory usage
462        let mut multi_index = MemoryMultiIndex::new();
463
464        // Process each attribute index entry, but only load the ones needed for our query
465        let query_fields: Vec<String> = query.iter().map(|(field, _, _)| field.clone()).collect();
466
467        for attr_info in attr_index_entries.iter() {
468            let column_idx = attr_info.index();
469            let field_name = columns[column_idx as usize].name().to_string();
470
471            // Only process this attribute if it's used in the query
472            if query_fields.contains(&field_name) {
473                add_indices_to_multi_memory_index(
474                    &mut self.reader,
475                    &mut multi_index,
476                    &columns,
477                    &query,
478                    attr_info,
479                )?;
480            } else {
481                // Skip this attribute index if not needed
482                let index_size = attr_info.length();
483                io::copy(
484                    &mut (&mut self.reader).take(index_size as u64),
485                    &mut io::sink(),
486                )?;
487            }
488        }
489
490        // Build and execute the query
491        let query_obj = build_query(&query);
492        let mut result = multi_index.query(&query_obj.conditions)?;
493        result.sort();
494
495        let header_size = self.buffer.header_buf.len();
496        let feature_offset = FeatureOffset {
497            magic_bytes: 8,
498            header: header_size as u64,
499            rtree_index: self.rtree_index_size(),
500            attributes: self.attr_index_size(),
501        };
502
503        let total_feat_count = result.len() as u64;
504
505        // Create and return the FeatureIter
506        Ok(FeatureIter::<R, NotSeekable>::new(
507            self.reader,
508            self.verify,
509            self.buffer,
510            None,
511            Some(result),
512            feature_offset,
513            total_feat_count,
514        ))
515    }
516}
fcb_core/reader/attr_query.rs

fcb_core/reader/
attr_query.rs
