polars-core 0.53.0

Core of the Polars DataFrame library
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

use polars_error::{PolarsResult, polars_err};
use polars_utils::aliases::PlHashMap;

use crate::frame::DataFrame;
use crate::prelude::Column;

pub(super) const LINEAR_SEARCH_LIMIT: usize = 4;

/// Selects columns by name.
pub(super) enum AmortizedColumnSelector<'a> {
    Direct {
        df: &'a DataFrame,
    },
    NameToIdxMapping {
        df: &'a DataFrame,
        name_to_idx: PlHashMap<&'a str, usize>,
    },
}

impl<'a> AmortizedColumnSelector<'a> {
    pub(super) fn new(df: &'a DataFrame) -> Self {
        if df.width() <= LINEAR_SEARCH_LIMIT || df.cached_schema().is_some() {
            Self::Direct { df }
        } else {
            Self::NameToIdxMapping {
                df,
                name_to_idx: df
                    .columns()
                    .iter()
                    .enumerate()
                    .map(|(i, s)| (s.name().as_str(), i))
                    .collect(),
            }
        }
    }

    fn select(&self, name: &str) -> PolarsResult<&'a Column> {
        match self {
            Self::Direct { df } => df.column(name),
            Self::NameToIdxMapping { df, name_to_idx } => {
                let i = *name_to_idx
                    .get(name)
                    .ok_or_else(|| polars_err!(col_not_found = name))?;

                Ok(df.select_at_idx(i).unwrap())
            },
        }
    }

    /// Does not error on duplicate selections.
    pub(super) fn select_multiple(
        &self,
        names: impl IntoIterator<Item = impl AsRef<str>>,
    ) -> PolarsResult<Vec<Column>> {
        names
            .into_iter()
            .map(|name| self.select(name.as_ref()).cloned())
            .collect()
    }
}