trs_dataframe/dataframe/
index.rs

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
use data_value::DataValue;

use super::{colums_store::ColumnFrame, Key};

/// Index for the fast search in the candidates
/// This is used for joining the dataframes by the given keys
#[derive(Debug)]
pub struct Index {
    /// indexes for the values in the [`ColumnFrame`]
    index: Vec<usize>,
    /// actual values under the keys
    values: Vec<Vec<DataValue>>,
}

impl Index {
    /// Create the index for the given keys and the [`ColumnFrame`] for the given keys.
    /// This will enumerate the values and store them in the index with current values
    pub fn new(key: Vec<Key>, df: &ColumnFrame) -> Self {
        let selected = df.select(Some(key.as_slice()));
        let mut this = Self {
            index: Vec::new(),
            values: Vec::new(),
        };

        for (index, candidate) in selected.rows().into_iter().enumerate() {
            this.index.push(index);
            this.values.push(candidate.to_vec());
        }
        this
    }

    pub fn get(&self, values: &[DataValue]) -> Option<usize> {
        self.values
            .iter()
            .position(|row| row.iter().zip(values.iter()).all(|(a, b)| a == b))
            .map(|idx| self.index[idx])
    }
}